From 658817a7a068bb4734adbb8a7dae488ecdc526b1 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 14 Oct 2026 09:02:29 -0700 Subject: [PATCH 01/14] [None][feat] MiniMax-M3 NVFP4 TP4 decode target for modeling_v2, with step-loop runtime hooks Adds the modeling_v2 target minimax_m3_nvfp4__sm_103__tp4 (GB300, TP4, Eagle3 one-model speculation) and the executor seams it opts into. - Runtime hooks (off unless a model declares them): packed step inputs, skipped unchanged H2D uploads, a private graph metadata pool, sampler greedy-token and step-fold paths, MSA runtime policy, and rank-merged autotuned tactics. Greedy-only outputs are kept only when nothing reads the logits (no guided decoder, non-empty batch). - TP4 decode target: fused attention with long-context selection, a pre-wait-staged NVFP4 QKV/indexer producer and index scorer, multicast all-reduce + RMSNorm boundaries with the MoE finalize folded in, MXFP8 GEMVs, CuTe DSL MoE engines, an exact greedy and sampled verify, a fused draft loop, and lossless 12-bit coded BF16 target and draft heads. - Explicit dispatch: every covered shape is a certified catalog cell; under modeling_v2: require a covered shape that cannot run raises. Every item is on by default with one constant that turns it off; no environment reads. - Numerics match stock up to summation order; the long-context selector breaks >1024-way score ties by its contract (smaller block). Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/CMakeLists.txt | 18 +- cpp/tensorrt_llm/kernels/minimaxM3/attnC1.cu | 3846 ++++++++++++++++ cpp/tensorrt_llm/kernels/minimaxM3/attnC1.h | 110 + .../kernels/minimaxM3/attnMega.cu | 3673 ++++++++++++++++ cpp/tensorrt_llm/kernels/minimaxM3/attnMega.h | 112 + .../kernels/minimaxM3/draftStep.cu | 576 +++ .../kernels/minimaxM3/draftStep.h | 99 + .../kernels/minimaxM3/fusedDecodeAttn.cu | 3679 ++++++++++++++++ .../kernels/minimaxM3/fusedDecodeAttn.h | 105 + .../kernels/minimaxM3/mcAllreduceNorm.cu | 563 +++ .../kernels/minimaxM3/mcAllreduceNorm.h | 76 + .../kernels/minimaxM3/msaKvUpdate.cu | 70 + .../kernels/minimaxM3/msaKvUpdate.h | 45 + .../kernels/minimaxM3/mxfp8LowMGemm.cu | 820 ++++ .../kernels/minimaxM3/mxfp8LowMGemm.h | 46 + .../kernels/minimaxM3/producerPw.cu | 324 ++ .../kernels/minimaxM3/producerPw.h | 69 + .../kernels/minimaxM3/routerGemv.cu | 606 +++ .../kernels/minimaxM3/routerGemv.h | 46 + .../kernels/minimaxM3/sampledAccept.cu | 525 +++ .../kernels/minimaxM3/sampledAccept.h | 79 + .../kernels/minimaxM3/selectLong.cu | 235 + .../kernels/minimaxM3/selectLong.h | 45 + .../kernels/minimaxM3/shardAccept.cu | 204 + .../kernels/minimaxM3/shardAccept.h | 67 + .../kernels/minimaxM3/sparseDecodeLc.cu | 561 +++ .../kernels/minimaxM3/sparseDecodeLc.h | 70 + .../kernels/minimaxM3/specTail.cu | 425 ++ cpp/tensorrt_llm/kernels/minimaxM3/specTail.h | 80 + .../kernels/minimaxM3/targetHead.cu | 372 ++ .../kernels/minimaxM3/targetHead.h | 68 + cpp/tensorrt_llm/kernels/stepInputsScatter.cu | 69 + cpp/tensorrt_llm/kernels/stepInputsScatter.h | 40 + cpp/tensorrt_llm/thop/CMakeLists.txt | 16 + cpp/tensorrt_llm/thop/minimaxM3AttnC1Op.cpp | 191 + cpp/tensorrt_llm/thop/minimaxM3AttnMegaOp.cpp | 207 + .../thop/minimaxM3DraftStepOp.cpp | 198 + .../thop/minimaxM3FusedDecodeAttnOp.cpp | 175 + .../thop/minimaxM3McAllreduceNormOp.cpp | 159 + .../thop/minimaxM3MsaKvUpdateOp.cpp | 110 + .../thop/minimaxM3Mxfp8LowMGemmOp.cpp | 83 + .../thop/minimaxM3ProducerPwOp.cpp | 181 + .../thop/minimaxM3RouterGemvOp.cpp | 149 + .../thop/minimaxM3SampledAcceptOp.cpp | 105 + .../thop/minimaxM3SelectLongOp.cpp | 80 + .../thop/minimaxM3ShardAcceptOp.cpp | 118 + .../thop/minimaxM3SparseDecodeLcOp.cpp | 139 + cpp/tensorrt_llm/thop/minimaxM3SpecTailOp.cpp | 202 + .../thop/minimaxM3TargetHeadOp.cpp | 123 + cpp/tensorrt_llm/thop/stepInputsScatterOp.cpp | 57 + .../_experimental/modeling_v2/__init__.py | 2 + .../modeling_v2/_router_index.py | 22 +- .../modeling_v2/models/minimax_m3/__init__.py | 3 + .../minimax_m3_nvfp4__sm_103__tp4/__init__.py | 3 + .../minimax_m3_nvfp4__sm_103__tp4/dispatch.py | 580 +++ .../kernels/__init__.py | 3 + .../kernels/long_select.py | 62 + .../kernels/m3_attn_lc.py | 112 + .../kernels/m3_attn_mega.py | 199 + .../kernels/m3_attn_mega_c1.py | 358 ++ .../kernels/m3_draft_step.py | 179 + .../kernels/m3_fused_attn.py | 349 ++ .../kernels/m3_index_score.py | 645 +++ .../kernels/m3_mc_allreduce.py | 368 ++ .../kernels/m3_prequant.py | 111 + .../kernels/m3_router_gemv.py | 107 + .../kernels/m3_sampled_accept.py | 158 + .../kernels/m3_shard_accept.py | 176 + .../kernels/m3_spec_tail.py | 113 + .../kernels/m3_target_head.py | 113 + .../kernels/msa_kv_update.py | 89 + .../kernels/mxfp8_low_m_gemm.py | 109 + .../kernels/triton_sparse_decode.py | 1651 +++++++ .../kernels/triton_sparse_decode_mq.py | 237 + .../minimax_m3_nvfp4__sm_103__tp4/layers.py | 1682 +++++++ .../minimax_m3_nvfp4__sm_103__tp4/linears.py | 259 ++ .../minimax_m3_nvfp4__sm_103__tp4/modeling.py | 633 +++ .../ops/__init__.py | 3 + .../ops/attention/__init__.py | 3 + .../ops/attention/m3_attn_c1.py | 280 ++ .../ops/attention/m3_attn_mega.py | 340 ++ .../ops/attention/m3_dense_decode.py | 187 + .../ops/attention/m3_fused_decode_attn.py | 643 +++ .../ops/attention/m3_index_score.py | 159 + .../ops/attention/m3_producer.py | 197 + .../ops/attention/m3_select_long.py | 172 + .../ops/attention/m3_sparse_decode.py | 215 + .../ops/attention/m3_sparse_decode_lc.py | 395 ++ .../ops/attention/msa_kv_update.py | 200 + .../ops/comm/__init__.py | 3 + .../ops/comm/m3_mc_allreduce_norm.py | 277 ++ .../ops/comm/m3_mc_finalize_allreduce_norm.py | 197 + .../ops/gemm/__init__.py | 3 + .../ops/gemm/m3_mx_gemv.py | 166 + .../ops/gemm/m3_mx_gemv_bf16.py | 106 + .../ops/gemm/m3_target_head.py | 346 ++ .../ops/gemm/mxfp8_low_m_gemm.py | 176 + .../ops/moe/__init__.py | 3 + .../ops/moe/m3_moe16_engine.py | 67 + .../ops/moe/m3_moe_engine.py | 325 ++ .../ops/moe/m3_router_gemv.py | 101 + .../ops/moe/m3_router_gemv_fp4.py | 63 + .../ops/moe/m3_router_tc.py | 106 + .../ops/norm/__init__.py | 3 + .../ops/norm/m3_embedding_firstnorm.py | 89 + .../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_step.py | 367 ++ .../ops/spec/m3_sampled_accept.py | 269 ++ .../ops/spec/m3_sampled_accept_ws_words.py | 43 + .../ops/spec/m3_shard_accept.py | 284 ++ .../ops/spec/m3_slot_scatter.py | 72 + .../ops/spec/m3_step_gather.py | 152 + .../minimax_m3_nvfp4__sm_103__tp4/runtime.py | 50 + .../spec_fold.py | 112 + .../spec_worker.py | 1166 +++++ .../minimax_m3_nvfp4__sm_103__tp4/weights.py | 62 + .../modeling_v2/models/minimax_m3/routing.py | 149 + .../attention/backends/fmha/msa_decode.py | 21 +- .../attention/backends/fmha/msa_prefill.py | 43 +- .../sparse/minimax_m3/cache_manager.py | 124 +- .../backends/sparse/minimax_m3/common.py | 29 +- .../backends/sparse/minimax_m3/msa_backend.py | 208 +- .../_torch/attention/backends/trtllm.py | 13 +- .../_torch/custom_ops/cute_dsl_custom_ops.py | 20 +- .../_torch/custom_ops/torch_custom_ops.py | 68 + .../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 | 1236 ++++++ .../cute_dsl_kernels/minimax_m3/mx_gemv/op.py | 474 ++ tensorrt_llm/_torch/modules/linear.py | 10 + tensorrt_llm/_torch/pyexecutor/_util.py | 8 +- .../_torch/pyexecutor/cuda_graph_runner.py | 44 +- .../engine/runners/decoder/runner.py | 247 +- tensorrt_llm/_torch/pyexecutor/h2d_cache.py | 137 + .../_torch/pyexecutor/model_engine.py | 5 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 56 +- .../_torch/pyexecutor/py_executor_creator.py | 1 + .../_torch/pyexecutor/runtime_hooks.py | 198 + .../_torch/pyexecutor/sampler/sampler.py | 62 +- tensorrt_llm/_torch/pyexecutor/step_inputs.py | 194 + tensorrt_llm/_torch/speculative/eagle3.py | 7 +- tensorrt_llm/_torch/speculative/interface.py | 7 + .../_torch/speculative/spec_sampler_base.py | 93 +- .../msa/test_minimax_m3_index_decode_score.py | 37 +- .../attention/sparse/msa/test_msa_backend.py | 120 + .../sparse/msa/test_msa_runtime_policy.py | 362 ++ .../engine/test_decoder_forward_scopes.py | 242 ++ .../engine/test_decoder_runtime_hooks.py | 104 + .../_torch/executor/test_cuda_graph_runner.py | 116 + .../_torch/executor/test_h2d_cache.py | 316 ++ .../_torch/executor/test_py_executor.py | 1 + .../_torch/executor/test_runtime_hooks.py | 111 + .../_torch/executor/test_step_inputs.py | 604 +++ .../attention/_m3_fused_attn_step.py | 255 ++ .../attention/_m3_nvfp4_cache.py | 104 + .../attention/test_modeling_v2_m3_attn_c1.py | 137 + .../test_modeling_v2_m3_attn_mega.py | 157 + .../test_modeling_v2_m3_dense_decode.py | 149 + .../test_modeling_v2_m3_fused_decode_attn.py | 129 + .../test_modeling_v2_m3_index_score.py | 92 + .../attention/test_modeling_v2_m3_producer.py | 179 + .../test_modeling_v2_m3_select_long.py | 159 + .../test_modeling_v2_m3_sparse_decode.py | 121 + .../test_modeling_v2_m3_sparse_decode_lc.py | 237 + .../test_modeling_v2_msa_kv_update.py | 128 + .../comm/_m3_draft_ar_norm_op_matrix.py | 380 ++ .../comm/_m3_mc_allreduce_norm_op_matrix.py | 663 +++ .../comm/_rank_job.py | 81 + .../comm/test_modeling_v2_m3_draft_ar_norm.py | 30 + ...eling_v2_m3_mc_allreduce_norm_op_matrix.py | 38 + .../gemm/test_modeling_v2_m3_mx_gemv.py | 86 + .../gemm/test_modeling_v2_m3_mx_gemv_bf16.py | 63 + .../gemm/test_modeling_v2_m3_target_head.py | 248 ++ .../gemm/test_modeling_v2_mxfp8_low_m_gemm.py | 99 + .../moe/_m3_moe_operands.py | 170 + .../moe/test_modeling_v2_m3_moe16_engine.py | 58 + .../moe/test_modeling_v2_m3_moe_engine.py | 113 + .../moe/test_modeling_v2_m3_router_gemv.py | 148 + ...test_modeling_v2_m3_embedding_firstnorm.py | 71 + .../spec/_m3_spec_inputs.py | 106 + .../spec/test_modeling_v2_m3_capture_norm.py | 57 + .../spec/test_modeling_v2_m3_chunk_norm.py | 50 + .../spec/test_modeling_v2_m3_draft_input.py | 98 + .../spec/test_modeling_v2_m3_draft_step.py | 686 +++ .../test_modeling_v2_m3_exact_numerics.py | 293 ++ .../test_modeling_v2_m3_sampled_accept.py | 167 + .../spec/test_modeling_v2_m3_shard_accept.py | 120 + .../spec/test_modeling_v2_m3_slot_scatter.py | 57 + .../spec/test_modeling_v2_m3_step_gather.py | 80 + .../test_modeling_v2_m3_capture_boundary.py | 65 + .../test_modeling_v2_m3_dispatch.py | 836 ++++ .../test_modeling_v2_m3_prebuild.py | 316 ++ .../test_modeling_v2_m3_rank_agreement.py | 190 + .../modeling_v2/test_modeling_v2_routing.py | 14 +- .../_torch/sampler/test_greedy_tokens.py | 205 + .../test_spec_sampler_step_fold.py | 254 ++ 207 files changed, 57801 insertions(+), 164 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/attnC1.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/attnC1.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/attnMega.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/attnMega.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/draftStep.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/draftStep.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/mcAllreduceNorm.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/mcAllreduceNorm.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/msaKvUpdate.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/msaKvUpdate.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/mxfp8LowMGemm.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/mxfp8LowMGemm.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/producerPw.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/producerPw.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/selectLong.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/selectLong.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/shardAccept.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/shardAccept.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/sparseDecodeLc.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/sparseDecodeLc.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/specTail.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/targetHead.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3/targetHead.h create mode 100644 cpp/tensorrt_llm/kernels/stepInputsScatter.cu create mode 100644 cpp/tensorrt_llm/kernels/stepInputsScatter.h create mode 100644 cpp/tensorrt_llm/thop/minimaxM3AttnC1Op.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3AttnMegaOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3DraftStepOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3FusedDecodeAttnOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3McAllreduceNormOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3MsaKvUpdateOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3Mxfp8LowMGemmOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3ProducerPwOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3RouterGemvOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3SampledAcceptOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3SelectLongOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3ShardAcceptOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3SparseDecodeLcOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3SpecTailOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3TargetHeadOp.cpp create mode 100644 cpp/tensorrt_llm/thop/stepInputsScatterOp.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/dispatch.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/long_select.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_lc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega_c1.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_draft_step.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_fused_attn.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_mc_allreduce.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_prequant.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_router_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_sampled_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_shard_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_spec_tail.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_target_head.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/msa_kv_update.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mxfp8_low_m_gemm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/triton_sparse_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/triton_sparse_decode_mq.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/layers.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/linears.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_attn_c1.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_attn_mega.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_dense_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_fused_decode_attn.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_producer.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_select_long.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode_lc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/msa_kv_update.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_mc_allreduce_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp4/ops/gemm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv_bf16.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_target_head.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/mxfp8_low_m_gemm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe16_engine.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe_engine.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_router_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_router_gemv_fp4.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_router_tc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/m3_embedding_firstnorm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_capture_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_chunk_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_draft_input.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_draft_step.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_sampled_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp4/ops/spec/m3_shard_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_slot_scatter.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_step_gather.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/runtime.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec_fold.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec_worker.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/weights.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/routing.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/pyexecutor/h2d_cache.py create mode 100644 tensorrt_llm/_torch/pyexecutor/runtime_hooks.py create mode 100644 tensorrt_llm/_torch/pyexecutor/step_inputs.py create mode 100644 tests/unittest/_torch/attention/sparse/msa/test_msa_runtime_policy.py create mode 100644 tests/unittest/_torch/executor/engine/test_decoder_forward_scopes.py create mode 100644 tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py create mode 100644 tests/unittest/_torch/executor/test_h2d_cache.py create mode 100644 tests/unittest/_torch/executor/test_runtime_hooks.py create mode 100644 tests/unittest/_torch/executor/test_step_inputs.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/_m3_fused_attn_step.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/_m3_nvfp4_cache.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_c1.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_mega.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_dense_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_fused_decode_attn.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_index_score.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_producer.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_select_long.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_sparse_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp4/attention/test_modeling_v2_msa_kv_update.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_draft_ar_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_rank_job.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/gemm/test_modeling_v2_m3_mx_gemv.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp4/gemm/test_modeling_v2_m3_target_head.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_mxfp8_low_m_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/_m3_moe_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_moe16_engine.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_moe_engine.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_gemv.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/_m3_spec_inputs.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_capture_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_chunk_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_input.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_step.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_sampled_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_shard_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_slot_scatter.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_step_gather.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_capture_boundary.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_dispatch.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_prebuild.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_rank_agreement.py create mode 100644 tests/unittest/_torch/sampler/test_greedy_tokens.py create mode 100644 tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index 0e7fc9390dc0..9cbf1480a879 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -77,6 +77,10 @@ list(FILTER SRC_CU EXCLUDE REGEX "fusedSampling/.*") # Marlin is built as its own architecture-scoped OBJECT library below. list(FILTER SRC_CPP EXCLUDE REGEX "marlin/.*") list(FILTER SRC_CU EXCLUDE REGEX "marlin/.*") +# The MiniMax-M3 decode attention kernels are built as their own +# architecture-scoped OBJECT library below. +list(FILTER SRC_CU EXCLUDE REGEX + "minimaxM3/(fusedDecodeAttn|attnMega|attnC1)\\.cu$") if(NOT ENABLE_MULTI_DEVICE) list(FILTER SRC_CU EXCLUDE REGEX "customAllReduceKernels*.*cu$") @@ -111,10 +115,22 @@ if(MARLIN_SRC) set_cuda_architectures(marlin_src 89 90 120f) endif() +# MiniMax-M3 decode attention: block-scaled tcgen05 MMAs (.scale_vec::1X) exist +# on the arch-specific sm_100a / sm_103a targets only, not on the sm_100f family +# target the rest of this library builds. Other targets compile kernels that +# trap. +add_library( + minimax_m3_attn_src OBJECT minimaxM3/fusedDecodeAttn.cu minimaxM3/attnMega.cu + minimaxM3/attnC1.cu) +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) + 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/attnC1.cu b/cpp/tensorrt_llm/kernels/minimaxM3/attnC1.cu new file mode 100644 index 000000000000..e375b6c994fd --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/attnC1.cu @@ -0,0 +1,3846 @@ +/* + * 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 attention + o_proj kernel (c1): fused decode attention and o_proj GEMV in one 96-CTA launch. + +#include "tensorrt_llm/kernels/minimaxM3/attnC1.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// The kernel uses tcgen05/TMEM (including block-scaled MMAs, which family targets lack), cluster DSMEM and E2M1 +// conversions: its body is compiled for the arch-specific sm_100a / sm_103a targets only (see the minimax_m3_attn_src +// library in kernels/CMakeLists.txt). Any other target compiles a kernel that traps, so a launch on a build without +// those targets fails loudly instead of returning untouched outputs. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030) \ + && (defined(__CUDA_ARCH_SPECIFIC__) || defined(__CUDA_ARCH_FEAT_SM100_ALL) || defined(__CUDA_ARCH_FEAT_SM103_ALL)) +#define MINIMAX_M3_ATTN_DEVICE_ENABLED 1 +#else +#define MINIMAX_M3_ATTN_DEVICE_ENABLED 0 +#endif + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr int kHeadDim = 128; +constexpr int kRotDim = 64; +constexpr int kPage = 128; +constexpr int kHq = 16; +constexpr int kColK = kHq; // packed head slots: [Q x16 | K | V | idxQ | idxK] +constexpr int kColV = kHq + 1; +constexpr int kColIdxK = kHq + 3; +constexpr int kBlocks = 16; // KV blocks per request == topk == cluster size +constexpr int kComputeWarps = 16; // QK / softmax / PV and the dequant (one 16-byte chunk per thread) +constexpr int kWarps = kComputeWarps + 4; // plus producer-only warps, so new K/V/index-K never queue behind a Q head +constexpr int kThreads = kWarps * 32; +constexpr int kComputeThreads = kComputeWarps * 32; +constexpr int kOPitch = 132; // floats per partial-O row +constexpr float kNegInf = -__builtin_huge_valf(); + +using Params = MinimaxM3AttnC1Params; + +// Optional per-CTA clock64 probes for development (-DM3FA_PROFILE and a prof tensor): [CTA][64] timestamps. +// -DM3FA_PROFILE_GT records %globaltimer (ns, comparable across SMs and kernels) instead of clock64. +#ifdef M3FA_PROFILE_GT +#define M3FA_PROFILE +#define M3FA_NOW(v) asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(v)) +#else +#define M3FA_NOW(v) v = clock64() +#endif +#ifdef M3FA_PROFILE +#define M3FA_PROF_T(i, thread) \ + do \ + { \ + if (p.prof != nullptr && threadIdx.x == (thread)) \ + { \ + unsigned long long now_; \ + M3FA_NOW(now_); \ + p.prof[blockIdx.x * 64 + (i)] = now_; \ + } \ + } while (0) +#else +#define M3FA_PROF_T(i, thread) \ + do \ + { \ + } while (0) +#endif +#define M3FA_PROF(i) M3FA_PROF_T(i, 0) +// clock64 probes (development): slot i of this CTA, written by every calling thread. +#ifdef M3FA_PROFILE +#define M3FA_CLK(i) \ + do \ + { \ + if (p.prof != nullptr) \ + { \ + p.prof[blockIdx.x * 64 + (i)] = clock64(); \ + } \ + } while (0) +#else +#define M3FA_CLK(i) \ + do \ + { \ + } while (0) +#endif + +#ifndef M3FA_MEGA_EARLYCLW +#define M3FA_MEGA_EARLYCLW 1 // pass the cluster barrier (peers' merge mbarriers initialized) early, not before the push +#endif +#ifndef M3FA_V15_EARLYV +#define M3FA_V15_EARLYV 0 // packed path: dequantize V before the PDL wait (switch; off by default) +#endif +#ifndef M3FA_MEGA_TRIGGER +// Megakernel PDL release: 0 o_proj CTAs after their activation (attention CTAs at exit); 1 as 0 with the attention +// CTAs after their KV reads; 2 every CTA right after its griddepcontrol.wait; 3 o_proj CTAs after their epilogue. +#define M3FA_MEGA_TRIGGER 3 +#endif +#ifndef M3FA_TRIGGER_AT +#define M3FA_TRIGGER_AT 2 // where the PDL dependents are released: 0 entry, 1 after the wait, 2 after the KV reads +#endif + +// Generic-proxy shared-memory writes (fp16 tiles, Q, P) become visible to the tensor core (async proxy). +#define M3FA_FENCE_PROXY_ASYNC() asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory") + +__device__ __forceinline__ void cp_async16(void* dst, void const* src) +{ + uint32_t const d = static_cast(__cvta_generic_to_shared(dst)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(d), "l"(src) : "memory"); +} + +__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__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +// One packed E2M1 byte -> half2 (low nibble in the low half). +__device__ __forceinline__ uint32_t e2m1x2_to_h2(uint32_t byte) +{ + uint32_t out; + asm("{\n .reg .b8 src;\n mov.b32 {src, _, _, _}, %1;\n cvt.rn.f16x2.e2m1x2 %0, src;\n}\n" : "=r"(out) : "r"(byte)); + return out; +} + +__device__ __forceinline__ uint32_t hmul2_u(uint32_t a, uint32_t b) +{ + __half2 r = __hmul2(*reinterpret_cast<__half2*>(&a), *reinterpret_cast<__half2*>(&b)); + return *reinterpret_cast(&r); +} + +__device__ __forceinline__ uint32_t e4m3_to_h2_splat(uint32_t byte) +{ + __half_raw const h = __nv_cvt_fp8_to_halfraw(static_cast<__nv_fp8_storage_t>(byte), __NV_E4M3); + return static_cast(h.x) | (static_cast(h.x) << 16); +} + +__device__ __forceinline__ uint32_t e4m3x2_to_h2(uint32_t two_bytes) +{ + __half2_raw const h = __nv_cvt_fp8x2_to_halfraw2(static_cast<__nv_fp8x2_storage_t>(two_bytes), __NV_E4M3); + return static_cast(h.x) | (static_cast(h.y) << 16); +} + +// ---- Exact copies of the C++ producer's helpers (fusedQKNormRopeKernel.cu / quantization.cuh) ---- +template +__device__ __forceinline__ uint32_t cvt_sf_group_shfl_mask() +{ + constexpr uint32_t groupSize = static_cast(NUM_THREADS_PER_SF); + constexpr uint32_t groupMask = 0xFFFFFFFFU >> (32U - groupSize); + uint32_t laneId = 0; + asm("mov.u32 %0, %%laneid;" : "=r"(laneId)); + return groupMask << (laneId & ~(groupSize - 1U)); +} + +__device__ __forceinline__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) +{ + 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"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]), "f"(array[4]), "f"(array[5]), "f"(array[6]), + "f"(array[7])); + return val; +} + +__device__ __forceinline__ uint16_t quantizeMinimaxM3Fp4x4( + float const (&elements)[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(elements[i]); + localMax = fmaxf(localMax, fabsf(__bfloat162float(rounded[i]))); + } + uint32_t const groupMask = cvt_sf_group_shfl_mask<4>(); + 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 quantValues[8] = {}; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + quantValues[i] = __bfloat162float(rounded[i]) * outputScale; + } + return static_cast(fp32_vec_to_e2m1(quantValues)); +} + +// RMSNorm (gemma 1+w) + NeoX RoPE over the first 64 dims, exactly as the producer: lane owns 4 dims. +__device__ __forceinline__ void norm_rope( + float (&e)[4], float const (&w)[4], float eps, float const (&cs)[4], float const (&sn)[4], int lane) +{ + float ss = 0.0F; + ss += e[0] * e[0] + e[1] * e[1]; + ss += e[2] * e[2] + e[3] * e[3]; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + ss = ss + __shfl_xor_sync(0xffffffffu, ss, mask, 32); + } + float const rmsReciprocal = rsqrtf(ss / static_cast(kHeadDim) + eps); +#pragma unroll + for (int i = 0; i < 4; ++i) + { + e[i] *= rmsReciprocal * (1.0F + w[i]); + } + __syncwarp(); + constexpr int kPairOffset = (kRotDim / 2) / 4; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + int const dim = lane * 4 + i; + float paired = __shfl_xor_sync(0xffffffffu, e[i], kPairOffset); + if (dim < kRotDim) + { + bool const firstHalf = dim < kRotDim / 2; + if (firstHalf) + { + paired = -paired; + } + e[i] = e[i] * cs[i] + paired * sn[i]; + } + } + __syncwarp(); +} + +__device__ __forceinline__ void load_w4(float (&w)[4], __nv_bfloat16 const* src, int lane) +{ + uint2 const raw = *reinterpret_cast(src + lane * 4); + float2 const a = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw.x)); + float2 const b = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw.y)); + w[0] = a.x; + w[1] = a.y; + w[2] = b.x; + w[3] = b.y; +} + +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ uint32_t mapa_shared(uint32_t addr, uint32_t rank) +{ + uint32_t r; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n" : "=r"(r) : "r"(addr), "r"(rank)); + return r; +} + +// QLEN query tokens per request; QG of them per cluster (QLEN / QG clusters per request); BPC KV blocks per CTA +// (cluster of kBlocks / BPC CTAs). +// Byte offset of element (r, k) in a UMMA no-swizzle K-major operand tile of 128 k-elements per row: 8x8 core +// matrices of 16-byte rows, 128 bytes apart along k (LBO) and 2048 bytes apart per 8 rows (SBO). The same formula +// with r = token, k = dim stores V so that it is the MN-major (dims contiguous) A operand of O^T = V^T P^T. +__device__ __forceinline__ int kmaj_off(int r, int k) +{ + return (r >> 3) * 2048 + (k >> 3) * 128 + (r & 7) * 16 + (k & 7) * 2; +} + +__device__ __forceinline__ uint64_t umma_desc(uint32_t addr, uint32_t lbo, uint32_t sbo) +{ + return static_cast((addr >> 4) & 0x3FFFu) | (static_cast((lbo >> 4) & 0x3FFFu) << 16) + | (static_cast((sbo >> 4) & 0x3FFFu) << 32) | (1ull << 46); +} + +// kind::f16 instruction descriptor: F32 accumulate, F16 operands, B K-major, M = 128. +__host__ __device__ constexpr uint32_t umma_idesc(int n, int a_mn_major, int b_mn_major) +{ + return (1u << 4) | (static_cast(a_mn_major) << 15) | (static_cast(b_mn_major) << 16) + | (static_cast(n >> 3) << 17) | (static_cast(128 >> 4) << 24); +} + +// 8-row version for the split softmax: the same lane-bit order (16, 8, 4, then 2 and 1 in full) as +// warp_reduce_scatter16, so every row is combined in the same association (fp add is commutative, not associative). +// After it lanes with bits 4..2 = (b4, b3, b2) hold row 4*b4 + 2*b3 + b2. +template +__device__ __forceinline__ float warp_reduce_scatter8(float (&a)[8], int lane) +{ +#pragma unroll + for (int step = 0; step < 3; ++step) + { + int const mask = 16 >> step; + int const half = 4 >> step; + bool const upper = (lane & mask) != 0; +#pragma unroll + for (int i = 0; i < half; ++i) + { + float const send = upper ? a[i] : a[i + half]; + float const keep = upper ? a[i + half] : a[i]; + float const recv = __shfl_xor_sync(0xffffffffu, send, mask); + a[i] = kMax ? fmaxf(keep, recv) : keep + recv; + } + } + float x = a[0]; + float const o2 = __shfl_xor_sync(0xffffffffu, x, 2); + x = kMax ? fmaxf(x, o2) : x + o2; + float const o1 = __shfl_xor_sync(0xffffffffu, x, 1); + return kMax ? fmaxf(x, o1) : x + o1; +} + +// Reduce 16 per-lane values over the warp, one value per row: after the reduce-scatter lane l holds row +// 8*b4 + 4*b3 + 2*b2 + b1 (b_k = bit k of l) reduced over all 32 lanes. 16 shuffles instead of 80. +template +__device__ __forceinline__ float warp_reduce_scatter16(float (&a)[16], int lane) +{ +#pragma unroll + for (int step = 0; step < 4; ++step) + { + int const mask = 16 >> step; + int const half = 8 >> step; + bool const upper = (lane & mask) != 0; +#pragma unroll + for (int i = 0; i < half; ++i) + { + float const send = upper ? a[i] : a[i + half]; + float const keep = upper ? a[i + half] : a[i]; + float const recv = __shfl_xor_sync(0xffffffffu, send, mask); + a[i] = kMax ? fmaxf(keep, recv) : keep + recv; + } + } + float const other = __shfl_xor_sync(0xffffffffu, a[0], 1); + return kMax ? fmaxf(a[0], other) : a[0] + other; +} + +__device__ __forceinline__ void umma_f16( + uint32_t tmem_d, uint64_t adesc, uint64_t bdesc, uint32_t idesc, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %4, 0;\n tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(idesc), "r"(accumulate) + : "memory"); +} + +__device__ __forceinline__ void umma_commit(uint32_t bar) +{ + asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];\n" ::"r"(bar) : "memory"); +} + +__device__ __forceinline__ void mbar_wait_cta(uint32_t bar, uint32_t parity) +{ + uint32_t ready = 0; +#ifdef M3FA_DBG_HANG + long long spins = 0; +#endif + while (!ready) + { +#ifdef M3FA_DBG_HANG + if (++spins == (1LL << 24)) + { + printf("M3FA hang: block %d thread %d mbar 0x%x parity %u\n", blockIdx.x, threadIdx.x, bar, parity); + } +#endif + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(bar), "r"(parity) + : "memory"); + } +} + +// Spin without suspending (v52's acc_done wait): a suspended try_wait wakes late on a tcgen05.commit arrival. +#ifndef M3FA_EPI_SPIN +#define M3FA_EPI_SPIN 1 +#endif +__device__ __forceinline__ void mbar_wait_epi(uint32_t bar, uint32_t parity) +{ + if (!M3FA_EPI_SPIN) + { + mbar_wait_cta(bar, parity); + return; + } + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(bar), "r"(parity) + : "memory"); + } +} + +__device__ __forceinline__ void tmem_ld8(uint32_t addr, float (&v)[8]) +{ + uint32_t r[8]; + asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]) + : "r"(addr)); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < 8; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +__device__ __forceinline__ void tmem_ld16(uint32_t addr, float (&v)[16]) +{ + uint32_t r[16]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, " + "%15}, [%16];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]) + : "r"(addr)); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < 16; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +// Warp-level fp16 MMA (m16n8k16, fp32 accumulate) and its shared-memory fragment loads. +__device__ __forceinline__ void ldsm_x4(uint32_t addr, uint32_t (&r)[4]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void ldsm_x2(uint32_t addr, uint32_t (&r)[2]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" : "=r"(r[0]), "=r"(r[1]) : "r"(addr)); +} + +__device__ __forceinline__ void ldsm_x2_trans(uint32_t addr, uint32_t (&r)[2]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0, %1}, [%2];\n" + : "=r"(r[0]), "=r"(r[1]) + : "r"(addr)); +} + +__device__ __forceinline__ void mma_f16_16816(float (&d)[4], uint32_t const (&a)[4], uint32_t const (&b)[2]) +{ + 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};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +// Register attention core (warp MMAs, no TMEM) for up to two tokens per cluster; the tcgen05 core above that. +// Off by default: on GB300 it measured ~0.25 us slower than the tcgen05 core at M=1 (GEMV + attention chain). +#ifndef M3FA_REG_CORE +#define M3FA_REG_CORE 0 +#endif +template +struct RegCore +{ + static constexpr bool kOn = M3FA_REG_CORE != 0 && QG <= 2; +}; + +template +struct Smem +{ + static constexpr int kRows = 16 * QG; // query rows (token, head) of this cluster = UMMA N + static constexpr int kTok = BPC * kPage; // block tokens per CTA + static constexpr int kTileBytes = kPage * kHeadDim * 2; // one fp16 K or V block tile in UMMA layout + static constexpr int kK16 = 0; + static constexpr int kV16 = kK16 + BPC * kTileBytes; + static constexpr int kQ = kV16 + BPC * kTileBytes; // Q rows, UMMA K-major [row][dim] + // P^T per block, UMMA MN-major [token][row], over the block's K tile: P is written after the QK MMAs (which + // read every K tile) have completed. + static constexpr int kP = kK16; + // Packed V tiles as copied from the cache: [block][128 rows x 64 data bytes | 1 KB of 4x4-quad scales]. + static constexpr int kRawV = kQ + kRows * kHeadDim * 2; + static constexpr int kRawVBlock = kPage * 64 + kPage * 8; + static constexpr int kRowScale = kRawV + BPC * kRawVBlock; + static constexpr int kRowLse = kRowScale + BPC * kRows * 4; + // Per-warp softmax partials [block][row][4 warps]: max, then sum. + static constexpr int kRedM = (kRowLse + BPC * kRows * 4 + 15) / 16 * 16; + static constexpr int kRedL = kRedM + BPC * kRows * 16; + // Partials pushed by every CTA of the cluster: [chunk][token][local head][kOPitch] floats, lse at column 128. + // When they do not fit beside the tiles they overlay the K/V tiles (and P, Q): peers then push only once every + // CTA of the cluster has completed its PV MMAs (a cluster barrier right before the push). + // Register core: per-warp row maxima and sums [block][row][warp of the block], two arrays. + static constexpr int kRedX = kRedL + BPC * kRows * 16; + static constexpr int kRedXBytes = RegCore::kOn ? 2 * BPC * kRows * (16 / BPC) * 4 : 0; + static constexpr int kStatsEnd = kRedX + kRedXBytes; + // Register core: P (fp16, [row][token] core matrices) at the start of its block's K tile, the scaled-free O + // partial ([row][kOPitch] fp32) behind it. + static constexpr int kObufOff = 8192; + static_assert(!RegCore::kOn || (kRows * kPage * 2 <= kObufOff && kObufOff + kRows * kOPitch * 4 <= kTileBytes), + "register-core P and O staging fit in a K tile"); + static constexpr int kRecvBytes = kBlocks * QG * BPC * kOPitch * 4; + static constexpr bool kRecvAlias = kStatsEnd + kRecvBytes + 32 > 227 * 1024; + static constexpr int kRecv = kRecvAlias ? 0 : kStatsEnd; + static constexpr int kMbar = kRecvAlias ? kStatsEnd : kStatsEnd + kRecvBytes; // merge, MMA mbarriers, TMEM addr + static constexpr int kBytes = kMbar + 32; + static constexpr int kTmemCols = BPC * 2 * kRows <= 32 ? 32 + : BPC * 2 * kRows <= 64 ? 64 + : BPC * 2 * kRows <= 128 ? 128 + : BPC * 2 * kRows <= 256 ? 256 + : 512; + static_assert(BPC * 2 * kRows <= 512, "TMEM columns"); + static_assert(!kRecvAlias || kRecvBytes <= kRawV, "aliased partials must fit in the tiles, Q and P"); + static_assert(kRows * kPage * 2 <= kTileBytes, "P^T fits in the K tiles"); + static_assert(kBytes <= 227 * 1024, "shared memory budget"); + static_assert(QLEN % QG == 0 && (BPC == 1 || BPC == 2), "unsupported geometry"); +}; + +// The packed NVFP4 bytes of one 16-byte chunk (32 dims of one block row) and its two block scales. +struct RawChunk +{ + uint4 data; + uint32_t scales; // byte 0: dims [0, 16), byte 1: dims [16, 32) of the chunk +}; + +// Chunk q of a block tile: lanes 8j..8j+7 of a warp take 8 consecutive rows of one 32-dim chunk, so every 16-byte +// fp16 store phase writes 8 distinct core-matrix rows (conflict-free in the UMMA layout). +__device__ __forceinline__ int chunk_row(int q) +{ + return (q >> 5) * 8 + (q & 7); +} + +__device__ __forceinline__ int chunk_col(int q) +{ + return (q >> 3) & 3; +} + +__device__ __forceinline__ RawChunk load_chunk( + uint8_t const* data_page, uint8_t const* scale_page, int row, int cc, bool is_v) +{ + RawChunk c; + c.data = *reinterpret_cast(data_page + row * 64 + cc * 16); + if (is_v) + { + // V scales: 4x4 token-quad order [row / 4][scale column][row % 4]. + uint8_t const* sc = scale_page + (row >> 2) * 32 + (row & 3); + c.scales = static_cast(sc[(cc * 2) * 4]) | (static_cast(sc[(cc * 2 + 1) * 4]) << 8); + } + else + { + c.scales = *reinterpret_cast(scale_page + row * 8 + cc * 2); + } + return c; +} + +__device__ __forceinline__ void store_chunk_fp16(uint8_t* tile, int row, int cc, RawChunk const& c, bool valid) +{ + uint32_t const s0 = e4m3_to_h2_splat(c.scales & 0xFFu); + uint32_t const s1 = e4m3_to_h2_splat((c.scales >> 8) & 0xFFu); + uint32_t const wd[4] = {c.data.x, c.data.y, c.data.z, c.data.w}; + uint32_t out[16]; +#pragma unroll + for (int i = 0; i < 16; ++i) + { + uint32_t const b = (wd[i >> 2] >> (8 * (i & 3))) & 0xFFu; + out[i] = valid ? hmul2_u(e2m1x2_to_h2(b), i < 8 ? s0 : s1) : 0u; + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + *reinterpret_cast(tile + kmaj_off(row, cc * 32 + j * 8)) + = make_uint4(out[4 * j], out[4 * j + 1], out[4 * j + 2], out[4 * j + 3]); + } +} + +// ---------------- Fused QKV projection: block-scaled MXFP8 on tcgen05 (N=2560, K=6144) ---------------- +// out[M, 2560] = mxfp8(x)[M, 6144] @ W^T, the computation of the CUTLASS MXFP8 GEMM: W is the A operand (128 output +// rows per MMA), the quantized activation the B operand (8 token columns), FP32 accumulation in TMEM. A unit is one +// (128-row tile, K split) pair; the 20 x kQkvSplits units run on the last CTAs of the grid. Before the PDL wait a +// CTA streams its unit's weight tiles and scale atoms through a shared-memory ring into TMEM; after it, it quantizes +// its x slice (mxfp8_quantize math), runs the MMAs and stores its FP32 partial rows. +// +// Hand-off without a flag: the partial buffers are double buffered by call parity and hold a sentinel (an FP32 NaN +// the MMA never produces) at rest, so a consumer polls the partials themselves. The producers of call n reset the +// other parity's buffers after their PDL wait; call n - 1, the last reader of those, is complete by then. +constexpr int kQkvN = 2560; +constexpr int kQkvK = 6144; +constexpr int kQkvKTiles = kQkvK / 128; +constexpr int kQkvRowTiles = kQkvN / 128; +#ifndef M3FA_QKV_SPLITS +#define M3FA_QKV_SPLITS 4 +#endif +#ifndef M3FA_QKV_RING +#define M3FA_QKV_RING kQkvMaxTiles // every tile of the unit in flight at once +#endif +#ifndef M3FA_QKV_NACC +#define M3FA_QKV_NACC 1 +#endif +#ifndef M3FA_QKV_SCALED_MMA +#define M3FA_QKV_SCALED_MMA 1 // 0: unscaled per-block MMAs with the scales applied in the epilogue (see below) +#endif +#ifndef M3FA_QKV_PIPE +// 1: the weight TMA is issued first at entry; the activation has its own region, copied right after the PDL wait; +// each weight tile's MMAs issue as soon as it lands (v52's order). 0: tiles into TMEM first, activation over the ring. +#define M3FA_QKV_PIPE 1 +#endif +#ifndef M3FA_QKV_A_SMEM +#define M3FA_QKV_A_SMEM 1 // scaled MMAs only; 1: they read the weight tiles from shared memory (no copy into TMEM) +#endif +// Default: block-scaled kind::mxf8f6f4 MMAs reading A (the weight tiles) in place from the ring, issued per tile as +// it lands (M3FA_QKV_PIPE), one accumulator; M3FA_QKV_A_SMEM=0 copies each tile into TMEM first. Alternative +// (M3FA_QKV_SCALED_MMA=0): one unscaled kind::f8f6f4 MMA per 32-wide K block into its own accumulator, A read in place +// from shared memory, the epilogue applying both UE8M0 scales. On GB300 an M128 x N8 x K32 MMA occupies the tensor pipe +// ~105 cycles block-scaled and ~60 unscaled, independent of N up to 64; the unscaled form's epilogue (320 accumulator +// columns) costs more than it saves. +constexpr bool kQkvScaled = M3FA_QKV_SCALED_MMA != 0; +constexpr bool kQkvASmem = !kQkvScaled || M3FA_QKV_A_SMEM != 0; +constexpr int kQkvSplits = M3FA_QKV_SPLITS; +constexpr int kQkvUnits = kQkvRowTiles * kQkvSplits; +constexpr int kQkvMaxTiles = (kQkvKTiles + kQkvSplits - 1) / kQkvSplits; +constexpr int kQkvRing = M3FA_QKV_RING; +constexpr int kQkvNacc = M3FA_QKV_NACC; +constexpr int kQkvGrid = 112; // 7 clusters of 16 (or 14 of 8): one wave with the attention +constexpr int kMaxTok = 8; +constexpr int kQkvWTile = 128 * 128; // E4M3 bytes of one 128 x 128 weight tile +constexpr int kQkvSfAtom = 512; // one 128-row x 4-block UE8M0 scale atom +constexpr uint32_t kQkvSentinel = 0xFFFFFFFFu; +constexpr int64_t kQkvPartStride = static_cast(kQkvSplits) * kMaxTok * kQkvN; // one parity +// TMEM columns of the one 512-column allocation of a CTA: attention [0, 64), then the GEMM's accumulators, scale +// atoms (A per tile, B per tile) and the A tiles (32 columns of 128 x 128 E4M3 each). +constexpr int kTmemAll = 512; +// Megakernel: the attention CTAs run no projection unit, so the GEMM starts at column 0 of a projection CTA. +constexpr int kTAcc = 0; +constexpr int kTSfa = kTAcc + 16; +constexpr int kTSfb = kTSfa + 4 * kQkvMaxTiles; +constexpr int kTA = (kTSfb + 4 * kQkvMaxTiles + 31) / 32 * 32; +static_assert(!kQkvScaled || kTA + 32 * kQkvMaxTiles <= kTmemAll, "TMEM columns"); +static_assert(kQkvScaled || kTAcc + 32 * kQkvMaxTiles <= kTmemAll, "TMEM columns (one accumulator per K block)"); +static_assert(!kQkvASmem || kQkvRing == kQkvMaxTiles, "shared-memory A needs every tile resident"); +static_assert(kQkvUnits <= kQkvGrid, "at most one unit per CTA"); +static_assert(kQkvNacc == 1 || kQkvNacc == 2 || kQkvNacc == 4, "accumulator count"); +static_assert(!M3FA_QKV_PIPE || kQkvScaled, "the pipelined projection uses the block-scaled MMAs"); +constexpr int kQkvOwnRows = 128 / kQkvSplits; // rows of a tile whose split-K sum a CTA owns +// Megakernel o_proj (v52 m3_mx_gemv for [6144, 2048]: split 2, contiguous k-tiles, one accumulator): unit +// o = CTA - 16 (the last 96 CTAs), row tile o / 2, split o % 2. +constexpr int kOpN = 6144; +constexpr int kOpK = 2048; +constexpr int kOpKTiles = kOpK / 128; +constexpr int kOpSplits = 2; +constexpr int kOpTiles = kOpKTiles / kOpSplits; +constexpr int kOpOwnRows = 128 / kOpSplits; +constexpr int kOpFirstCta = kQkvGrid - kOpN / 128 * kOpSplits; // 16 +constexpr int kOpSfBytes = kOpKTiles * 512; // one parity of the MXFP8 hand-off's scales +#ifndef M3FA_OP_L2PF +#define M3FA_OP_L2PF \ + 0 // 1: QKV CTAs prefetch their o_proj weight slice into L2 before the PDL wait (competes + // with the QKV weight stream for HBM behind the MoE) +#endif +// o_proj TMEM columns (after the QKV unit's MMAs completed): accumulator, A / B scale atoms, A (8 tiles x 32) +// c1 (attention + o_proj kernel): the attention holds columns [0, 32) while the o_proj weights stream in, so the +// o_proj accumulator and scale atoms sit above them. +constexpr int kOAcc = 32; +constexpr int kOSfa = 48; +constexpr int kOSfb = kOSfa + 4 * kOpTiles; +constexpr int kOA = 128; +static_assert(kOA + 32 * kOpTiles <= kTmemAll && kOSfb + 4 * kOpTiles <= kOA, "o_proj TMEM columns"); +constexpr int kQkvMbStride = kMaxTok + 4; // mailbox row pitch (fp32 words) + +__host__ __device__ constexpr int qkv_tile0(int s) +{ + return s * kQkvKTiles / kQkvSplits; +} + +__device__ __forceinline__ int qkv_unit() +{ + int const u = kQkvGrid - 1 - static_cast(blockIdx.x); + return u < kQkvUnits ? u : -1; +} + +// Attention layout plus, for the fused projection, the weight ring in front (the attention tiles reuse that space +// once the ring has drained into TMEM); the activation tiles, the pre-wait V copy and the barriers sit behind. +template +struct SmemX : Smem +{ + using B = Smem; + static constexpr int kRingW = 0; // [ring][128 rows][128 B], 128-byte swizzle + static constexpr int kRingSf = kRingW + kQkvRing * kQkvWTile; // [ring][512 B] + static constexpr int kRingEnd = kRingSf + kQkvRing * kQkvSfAtom; + // The activation tiles overlay the ring: written after the PDL wait, once every weight tile went on into TMEM. + static constexpr int kXt = kRingW; // [tile][8][128 B] + static constexpr int kXsf = kXt + kQkvMaxTiles * 1024; // [tile][512 B] + static constexpr int kXsc = (kRingEnd + 127) / 128 * 128; + // Peers push into recv at any time (an empty CTA right away), so it stays clear of the ring. + static constexpr int kRecv = FQKV ? kXsc : B::kRecv; + static constexpr int kRawV = FQKV ? (kRecv + kBlocks * QG * BPC * kOPitch * 4 + 127) / 128 * 128 : B::kRawV; + // Pipelined projection: the activation tiles and scale atoms in their own region (a projection CTA runs no + // attention, so it overlays the attention's partial-O and V regions), 1024-byte aligned for the swizzle. + static constexpr int kXA = FQKV ? (kRecv + 1023) / 1024 * 1024 : 0; + static constexpr int kXASf = kXA + kQkvMaxTiles * 1024; + static constexpr int kXAEnd = kXASf + kQkvMaxTiles * kQkvSfAtom; + static constexpr int kMbarA = FQKV ? kRawV + BPC * B::kRawVBlock : B::kMbar; + static constexpr int kMbar = FQKV && M3FA_QKV_PIPE && kXAEnd > kMbarA ? (kXAEnd + 127) / 128 * 128 : kMbarA; + static constexpr int kGBar = kMbar + 32; // full[ring], empty[ring], a_ready, acc_done, tile_done[tiles] + static constexpr int kRedBar = kGBar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); // split-K partials in + static constexpr int kOBar = kRedBar + 8; // o_proj: full[8], a_ready, acc_done, reduce + static constexpr int kMailbox = (kOBar + 8 * (kOpTiles + 3) + 15) / 16 * 16; // [split][own row][pitch] fp32 + // o_proj weight ring = ring tiles [0, 8); its activation tiles, scale atoms and mailbox behind it. + static constexpr int kOXt = kRingW + kOpTiles * kQkvWTile; + static constexpr int kOXsf = kOXt + kOpTiles * 1024; + static constexpr int kOMailbox = kOXsf + kOpTiles * kQkvSfAtom; + static_assert(!FQKV || kOMailbox + kOpSplits * kOpOwnRows * kQkvMbStride * 4 <= kRingW + kQkvRing * kQkvWTile, + "o_proj activation and mailbox fit behind its ring"); + static constexpr int kBytes = FQKV ? kMailbox + kQkvSplits * kQkvOwnRows * kQkvMbStride * 4 : B::kBytes; + static constexpr int kTmemCols = FQKV ? kTmemAll : B::kTmemCols; + static constexpr bool kRecvAlias = !FQKV && B::kRecvAlias; + // The epilogue reads the weight scale atoms while the attention already writes its tiles. + static_assert(!FQKV || kQkvScaled || B::kRecv <= kRingSf, "attention regions overlap the scale atoms"); + // The other attention regions overlay the ring: a unit's CTA writes them only once the ring has drained. + static_assert(kBytes <= 227 * 1024, "shared memory budget"); +}; + +// c1: the packed-input attention plus the o_proj unit blockIdx.x (tile bx / 2, split bx % 2). The o_proj weight ring +// and its barriers sit behind the attention layout; the activation tiles, scale atoms and the pair's split-K mailbox +// overlay the K tiles, free once this CTA has pushed its partials (a peer's mailbox push follows that push: its +// activation needs the merges of this cluster, which need this CTA's partials). +template +struct SmemC1 : SmemX +{ + using B = Smem; + static constexpr int kRingW = (B::kBytes + 1023) / 1024 * 1024; + static constexpr int kRingSf = kRingW + kOpTiles * kQkvWTile; + static constexpr int kOBar = kRingSf + kOpTiles * kQkvSfAtom; // full[8], a_ready, acc_done, reduce + static constexpr int kOXt = B::kK16; + static constexpr int kOXsf = kOXt + kOpTiles * 1024; + static constexpr int kOMailbox = kOXsf + kOpTiles * kQkvSfAtom; + static constexpr int kBytes = kOBar + 8 * (kOpTiles + 3); + static constexpr int kTmemCols = kTmemAll; + static_assert(!B::kRecvAlias, "c1: the partials have their own region"); + static_assert( + kOMailbox + kOpSplits * kOpOwnRows * kQkvMbStride * 4 <= B::kV16, "c1: o_proj overlay in the K tiles"); + static_assert(BPC * 2 * B::kRows <= kOAcc, "c1: attention TMEM columns below the o_proj's"); + static_assert(kBytes <= 227 * 1024, "c1: shared memory budget"); +}; + +__device__ __forceinline__ void mbar_init(uint32_t bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(bar), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbar_expect_tx(uint32_t bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(bar), "r"(bytes) : "memory"); +} + +__device__ __forceinline__ void bulk_g2s(void* dst, void const* src, uint32_t bytes, uint32_t bar) +{ + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], " + "%4;\n" ::"r"(smem_u32(dst)), + "l"(src), "r"(bytes), "r"(bar), "l"(0x12F0000000000000ull) // EVICT_FIRST: weights are read once per layer + : "memory"); +} + +__device__ __forceinline__ void tma_load_2d(uint32_t dst, void const* map, int c0, int c1, uint32_t bar) +{ + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;\n" ::"r"(dst), + "l"(map), "r"(c0), "r"(c1), "r"(bar), "l"(0x12F0000000000000ull) + : "memory"); +} + +// K-major operand in the 128-byte swizzle: rows of 128 bytes, 8-row groups 1024 bytes apart. +__device__ __forceinline__ uint64_t umma_desc_sw128(uint32_t addr) +{ + return umma_desc(addr, 16, 1024) | (2ull << 61); +} + +__device__ __forceinline__ void tmem_cp_128x256b(uint32_t taddr, uint64_t desc) +{ + asm volatile("tcgen05.cp.cta_group::1.128x256b [%0], %1;\n" ::"r"(taddr), "l"(desc)); +} + +__device__ __forceinline__ void tmem_cp_sf(uint32_t taddr, uint32_t saddr) +{ + asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;\n" ::"r"(taddr), "l"(umma_desc(saddr, 16, 128))); +} + +// Block-scaled MXFP8 instruction descriptor: E4M3 x E4M3, UE8M0 scales, M = 128, N = 8, scale byte kb. +__host__ __device__ constexpr uint32_t qkv_idesc(int kb) +{ + return (static_cast(kb) << 4) | ((8u >> 3) << 17) | (1u << 23) | ((128u >> 7) << 27) + | (static_cast(kb) << 29); +} + +__device__ __forceinline__ void umma_mx( + uint32_t tmem_d, uint32_t tmem_a, uint64_t bdesc, uint32_t idesc, uint32_t sfa, uint32_t sfb, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" + " tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [%0], [%1], %2, %3, [%4], [%5], p;\n}\n" :: + "r"(tmem_d), + "r"(tmem_a), "l"(bdesc), "r"(idesc), "r"(sfa), "r"(sfb), "r"(accumulate)); +} + +__device__ __forceinline__ float pow2_from_ue8m0(uint32_t e) +{ + return __uint_as_float(e == 0 ? 0x00400000u : (e << 23)); +} + +// Unscaled E4M3 x E4M3 MMA, FP32 accumulator: M = 128, N = 8, K = 32, a fresh accumulator. +constexpr uint32_t kQkvIdescF8 = (1u << 4) | ((8u >> 3) << 17) | ((128u >> 4) << 24); + +__device__ __forceinline__ void umma_f8_ss(uint32_t tmem_d, uint64_t adesc, uint64_t bdesc) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, 0, 0;\n tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(kQkvIdescF8)); +} + +__device__ __forceinline__ void tmem_ld16_nowait(uint32_t addr, float* v) +{ + uint32_t r[16]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, " + "%15}, [%16];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]) + : "r"(addr)); +#pragma unroll + for (int i = 0; i < 16; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +__device__ __forceinline__ void umma_mx_ss( + uint32_t tmem_d, uint64_t adesc, uint64_t bdesc, uint32_t idesc, uint32_t sfa, uint32_t sfb, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" + " tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [%0], %1, %2, %3, [%4], [%5], p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(idesc), "r"(sfa), "r"(sfb), "r"(accumulate)); +} + +__device__ __forceinline__ bool elect_one() +{ + uint32_t pred; + asm volatile("{\n .reg .pred P;\n elect.sync _|P, 0xffffffff;\n selp.u32 %0, 1, 0, P;\n}\n" : "=r"(pred)); + return pred != 0; +} + +__device__ __forceinline__ int qkv_parity(Params const& p) +{ + int v; + asm volatile("ld.relaxed.gpu.global.s32 %0, [%1];\n" : "=r"(v) : "l"(p.counters) : "memory"); + return v & 1; +} + +// Barriers and TMEM (warp 1), then the first ring stages (warp 0). Warps 0 and 1 of every CTA call it. +template +__device__ __forceinline__ void qkv_setup( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int lane, int u) +{ + constexpr bool kPipe = M3FA_QKV_PIPE != 0; + if (kPipe && warp == 0 && lane == 0 && u >= 0) + { + // First thing at entry: this unit's whole weight slice. The issuing thread initializes the ring's full + // barriers itself; every waiter on them passes bar.sync 8 below, which this thread reaches after the init. + asm volatile("prefetch.tensormap [%0];\n" ::"l"(&p.wmap) : "memory"); + int const rt = u / kQkvSplits; + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + for (int i = 0; i < kQkvRing; ++i) + { + mbar_init(gbar + 8 * i, 1); + } + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + for (int i = 0; i < kQkvRing && i < nt; ++i) + { + uint32_t const full = gbar + 8 * i; + mbar_expect_tx(full, kQkvWTile + kQkvSfAtom); + if (p.wpre != nullptr) + { + // Pre-tiled weight: the tile's 16 KB shared-memory image is contiguous (one bulk copy). + bulk_g2s(smem + SX::kRingW + i * kQkvWTile, + p.wpre + (static_cast(rt) * kQkvKTiles + kt0 + i) * kQkvWTile, kQkvWTile, full); + } + else + { + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap, (kt0 + i) * 128, rt * 128, full); + } + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf + (static_cast(rt) * kQkvKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, full); + } + M3FA_PROF_T(34, 0); + } + if (warp == 1) + { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n" ::"r"(smem_u32(tmem_slot)), + "r"(kTmemAll) + : "memory"); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n" ::: "memory"); + M3FA_PROF_T(35, 32); + if (lane == 0) + { + for (int i = (kPipe && u >= 0) ? kQkvRing : 0; i < 2 * kQkvRing + 2 + kQkvMaxTiles; ++i) + { + // PIPE 2: tile_done[0] (unused by the scaled MMAs) is act_ready, one arrival per thread outside warp 1 + mbar_init(gbar + 8 * i, M3FA_QKV_PIPE == 2 && i == 2 * kQkvRing + 2 ? kThreads - 32 : 1); + } + // The split-K reduce barrier: this CTA's own arrival carries the peers' partial bytes. + uint32_t const red = gbar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); + mbar_init(red, 1); + mbar_expect_tx(red, (kQkvSplits - 1) * kQkvOwnRows * kMaxTok * 4); + uint32_t const ob = smem_u32(smem + SX::kOBar); + for (int i = 0; i < kOpTiles + 3; ++i) + { + mbar_init(ob + 8 * i, 1); + } + mbar_expect_tx(ob + 8 * (kOpTiles + 2), (kOpSplits - 1) * kOpOwnRows * kMaxTok * 4); + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + asm volatile("bar.sync 8, 64;\n" ::: "memory"); + if (!kPipe && warp == 0 && lane == 0 && u >= 0) + { + int const rt = u / kQkvSplits; + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + for (int i = 0; i < kQkvRing && i < nt; ++i) + { + uint32_t const full = gbar + 8 * i; + mbar_expect_tx(full, kQkvWTile + kQkvSfAtom); + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap, (kt0 + i) * 128, rt * 128, full); + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf + (static_cast(rt) * kQkvKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, full); + } + M3FA_PROF_T(34, 0); + } +} + +// The rest of the ring (warp 0, lane 0) and the copies into TMEM (warp 1, lane 0); before the PDL wait. +template +__device__ __forceinline__ void qkv_stream_lane( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int u) +{ + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + uint32_t const full = gbar; + uint32_t const a_ready = gbar + 16 * kQkvRing; + if (warp == 0 || !kQkvScaled) + { + return; // every tile was issued in qkv_setup; the unscaled MMAs read them in place + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + // Fully unrolled with compile-time offsets: tcgen05 operands stay in uniform registers (a per-instruction + // conversion of computed operands costs ~150 cycles per MMA or copy). + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint32_t const sf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + mbar_wait_cta(full + 8 * i, 0); + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if constexpr (!kQkvASmem) + { +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kTA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + } + tmem_cp_sf(tm + kTSfa + i * 4, sf0 + i * kQkvSfAtom); + } + } + M3FA_PROF_T(31, 32); + umma_commit(a_ready); +} + +__device__ __forceinline__ bool mbar_test(uint32_t bar, uint32_t parity) +{ + uint32_t ready; + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(bar), "r"(parity) + : "memory"); + return ready != 0; +} + +// PIPE 2, warp 1's elected thread from before the PDL wait to the last MMA. Tiles that land before the activation +// is ready go on into TMEM (A and scale atom), off the critical path; once the activation is in (act_ready), the +// MMAs issue in tile order: the copied tiles with A in TMEM, every later tile as it lands with A read in place from +// the ring. The MMA sequence (and so the accumulation order) is the same either way. +template +__device__ __forceinline__ void qkv_hybrid(Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int u) +{ + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + uint32_t const act_ready = gbar + 8 * (2 * kQkvRing + 2); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + SX::kXA)); + uint32_t const wsf0 = smem_u32(smem + SX::kRingSf); + uint32_t const xsf0 = smem_u32(smem + SX::kXASf); + int copied = 0; // tiles [0, copied) are in TMEM + bool act = false; +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt && !act) + { + while (!mbar_test(gbar + 8 * i, 0)) + { + if (mbar_test(act_ready, 0)) + { + act = true; + break; + } + } + if (!act) + { + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kTA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + copied = i + 1; + } + } + } + while (!act) + { + act = mbar_test(act_ready, 0); + } + M3FA_PROF_T(21, 32); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + bool const in_tmem = i < copied; + if (!in_tmem) + { + mbar_wait_epi(gbar + 8 * i, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + } + uint32_t const sfb = tm + kTSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + uint64_t const bd = b0 + ((i * 1024 + kb * 32) >> 4); + if (in_tmem) + { + umma_mx(tm + kTAcc, tm + kTA + j * 8, bd, qkv_idesc(kb), tm + kTSfa + i * 4, sfb, j > 0); + } + else + { + umma_mx_ss(tm + kTAcc, w0 + ((i * kQkvWTile + kb * 32) >> 4), bd, qkv_idesc(kb), tm + kTSfa + i * 4, + sfb, j > 0); + } + } + } + } + M3FA_PROF_T(22, 32); + M3FA_PROF_T(31, 32); + umma_commit(gbar + 16 * kQkvRing); // a_ready (for its waiters) + umma_commit(gbar + 16 * kQkvRing + 8); // acc_done +} + +// The other lanes of warps 0 and 1 hold at __syncwarp rather than at griddepcontrol.wait, which would stall the +// whole warp (lane 0's loop included) until the producer grid completes. +template +__device__ __forceinline__ void qkv_stream( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int lane, int u) +{ + if (M3FA_QKV_PIPE == 2 && u >= 0 && warp == 1) + { + if (elect_one()) + { + qkv_hybrid(p, smem, gbar, tmem_slot, u); + } + __syncwarp(); + return; + } + if (u < 0 || warp > 1 || M3FA_QKV_PIPE) + { + return; // pipelined: the MMA thread takes each tile as it lands + } + if (lane == 0) + { + qkv_stream_lane(p, smem, gbar, tmem_slot, warp, u); + } + __syncwarp(); +} + +// After the PDL wait. Every thread quantizes a piece of the x slice into the B tiles (named barrier 9); warp 1 then +// issues the scale copies and the MMAs. +template +__device__ __forceinline__ void qkv_mma( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int tid, int u) +{ + int const warp = tid >> 5; + int const s = u % kQkvSplits; + int const kt0 = qkv_tile0(s); + int const nt = qkv_tile0(s + 1) - kt0; + { + // The MXFP8 activation slice (E4M3 rows of the unit's k-tiles in 128-byte-swizzled B tiles, and their scale + // atoms), copied with plain 16-byte loads over the overlaid ring once its tiles went on into TMEM: every load + // first, then the stores (one L2 round trip). + constexpr int kXDst = M3FA_QKV_PIPE ? SX::kXA : SX::kXt; + constexpr int kXSfDst = M3FA_QKV_PIPE ? SX::kXASf : SX::kXsf; + if (!M3FA_QKV_PIPE) + { + mbar_wait_cta(gbar + 16 * kQkvRing, 0); // a_ready: the ring's weight tiles are in TMEM + } + constexpr int kXItems = kQkvMaxTiles * 8 * kMaxTok; // (tile, 16-byte chunk, token) + constexpr int kSfItems = kQkvMaxTiles * (kQkvSfAtom / 16); + constexpr bool kHyb = M3FA_QKV_PIPE == 2; + constexpr int kCopiers = kHyb ? kThreads - 32 : kThreads; // PIPE 2: warp 1 is issuing the tiles + if (kHyb && warp == 1) + { + return; + } + int const ctid = kHyb && tid >= 64 ? tid - 32 : tid; + constexpr int kPer = (kXItems + kSfItems + kCopiers - 1) / kCopiers; + uint4 v[kPer]; + int dst[kPer]; +#pragma unroll + for (int k = 0; k < kPer; ++k) + { + int const q = ctid + kCopiers * k; + dst[k] = -1; + if (q < kXItems) + { + int const t = q % kMaxTok; + int const c = (q / kMaxTok) % 8; + int const i = q / (kMaxTok * 8); + if (i < nt && t < p.num_tokens) + { + v[k] = *reinterpret_cast( + p.xq + static_cast(t) * p.ldxq + (kt0 + i) * 128 + c * 16); + dst[k] = kXDst + i * 1024 + t * 128 + ((c ^ t) * 16); + } + } + else if (q < kXItems + kSfItems) + { + int const j = q - kXItems; + int const i = j / (kQkvSfAtom / 16); + int const c = j % (kQkvSfAtom / 16); + if (i < nt) + { + v[k] = *reinterpret_cast( + p.xsf_act + static_cast(kt0 + i) * kQkvSfAtom + c * 16); + dst[k] = kXSfDst + i * kQkvSfAtom + c * 16; + } + } + } +#pragma unroll + for (int k = 0; k < kPer; ++k) + { + if (dst[k] >= 0) + { + *reinterpret_cast(smem + dst[k]) = v[k]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if (kHyb) + { + asm volatile("mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];\n" ::"r"(gbar + 8 * (2 * kQkvRing + 2)) + : "memory"); + return; + } + if (warp != 1) + { + asm volatile("bar.arrive 9, %0;\n" ::"n"(kThreads) : "memory"); + } + } + if (warp == 1) + { + M3FA_PROF_T(32, 32); + asm volatile("bar.sync 9, %0;\n" ::"n"(kThreads) : "memory"); + M3FA_PROF_T(21, 32); + // Operands computed by the whole warp (warp-uniform, kept in uniform registers); one elected lane issues. + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + (M3FA_QKV_PIPE ? SX::kXA : SX::kXt))); + uint64_t const a0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + if (elect_one()) + { + M3FA_CLK(41); + if constexpr (!kQkvScaled) + { +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + mbar_wait_cta(gbar + 8 * i, 0); // weight tile i landed + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + umma_f8_ss(tm + kTAcc + (i * 4 + kb) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4), + b0 + ((i * 1024 + kb * 32) >> 4)); + } + umma_commit(gbar + 8 * (2 * kQkvRing + 2 + i)); // tile i's blocks are in TMEM + } + } + } + else + { + uint32_t const xsf0 = smem_u32(smem + (M3FA_QKV_PIPE ? SX::kXASf : SX::kXsf)); + uint32_t const wsf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + if constexpr (M3FA_QKV_PIPE != 0) + { + // Tile i as it lands: its A copy and scale atom into TMEM, then its MMAs (issue order). + mbar_wait_cta(gbar + 8 * i, 0); + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + if (i == nt - 1) + { + M3FA_CLK(47); // the unit's last weight tile landed (clock64 builds) + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if constexpr (!kQkvASmem) + { +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b( + tm + kTA + (i * 4 + kb) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + } + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + } + uint32_t const sfb = tm + kTSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + uint64_t const bd = b0 + ((i * 1024 + kb * 32) >> 4); + if constexpr (kQkvASmem) + { + umma_mx_ss(tm + kTAcc + (j % kQkvNacc) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4), bd, + qkv_idesc(kb), tm + kTSfa + i * 4, sfb, j >= kQkvNacc); + } + else + { + umma_mx(tm + kTAcc + (j % kQkvNacc) * 8, tm + kTA + j * 8, bd, qkv_idesc(kb), + tm + kTSfa + i * 4, sfb, j >= kQkvNacc); + } + } + } + } + } + M3FA_CLK(42); + if (M3FA_QKV_PIPE) + { + M3FA_PROF_T(31, 32); + umma_commit(gbar + 16 * kQkvRing); // a_ready (kept for its waiters: every tile is consumed) + } + umma_commit(gbar + 16 * kQkvRing + 8); // acc_done + } + M3FA_PROF_T(22, 32); + } +} + +// Warps 16..19 (TMEM lanes 32 (warp % 4) + lane = the unit's rows): partial rows out, the other parity reset. +template +__device__ __forceinline__ void qkv_epilogue( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int tid, int u, int par, unsigned int calls) +{ + int const warp = tid >> 5; + int const lane = tid & 31; + if constexpr (kQkvScaled) + { + mbar_wait_epi(gbar + 16 * kQkvRing + 8, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + M3FA_PROF_T(24, 512); + if (tid == 512) + { + M3FA_CLK(43); + } + int const r = (warp & 3) * 32 + lane; // row within the unit's tile + uint32_t const taddr = *tmem_slot + (((warp & 3) * 32) << 16) + kTAcc; + float val[kMaxTok]; + if constexpr (kQkvScaled) + { + float v[32]; + tmem_ld16(taddr, *reinterpret_cast(v)); + if constexpr (kQkvNacc == 4) + { + tmem_ld16(taddr + 16, *reinterpret_cast(v + 16)); + } +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = kQkvNacc == 4 ? (v[t] + v[8 + t]) + (v[16 + t] + v[24 + t]) + : kQkvNacc == 2 ? v[t] + v[8 + t] + : v[t]; + } + } + else + { + // Sum over K blocks in order: block result x (weight scale x activation scale). + int const nt = qkv_tile0(u % kQkvSplits + 1) - qkv_tile0(u % kQkvSplits); + float const* xsc = reinterpret_cast(smem + SX::kXsc); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = 0.0f; + } +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + // Tile i's four blocks as soon as its MMAs complete: the epilogue trails the MMAs by one tile. + mbar_wait_cta(gbar + 8 * (2 * kQkvRing + 2 + i), 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + float d[32]; + tmem_ld16_nowait(taddr + i * 32, d); + tmem_ld16_nowait(taddr + i * 32 + 16, d + 16); + uint32_t const wsw = *reinterpret_cast( + smem + SX::kRingSf + i * kQkvSfAtom + (r % 32) * 16 + (r / 32) * 4); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + float const ws = pow2_from_ue8m0((wsw >> (8 * kb)) & 0xFFu); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = fmaf(d[kb * 8 + t], ws * xsc[(i * 4 + kb) * kMaxTok + t], val[t]); + } + } + } + } + } + // v52's split-K reduce: split s owns rows [32 s, 32 s + 32) of the tile; the other splits push their fp32 partials + // of those rows into its mailbox (st.async, completing its reduce barrier's bytes); the owner adds the splits in + // rank order from 0.0 and stores bf16. + int const s = u % kQkvSplits; + int const owner = r / kQkvOwnRows; + int const local = r % kQkvOwnRows; + uint32_t const red = gbar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); + uint32_t cta_rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + // Units are numbered down from the last CTA: split q of this tile is CTA rank cta_rank - (q - s). + if (owner != s) + { + uint32_t const peer = cta_rank + s - owner; + uint32_t const dst + = mapa_shared(smem_u32(smem + SX::kMailbox) + ((s * kQkvOwnRows + local) * kQkvMbStride) * 4, peer); + uint32_t const pbar = mapa_shared(red, peer); +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];\n" ::"r"( + dst + 16 * g), + "r"(__float_as_uint(val[4 * g])), "r"(__float_as_uint(val[4 * g + 1])), + "r"(__float_as_uint(val[4 * g + 2])), "r"(__float_as_uint(val[4 * g + 3])), "r"(pbar) + : "memory"); + } + } + else + { + M3FA_PROF_T(26, tid - lane); + // test_wait (no suspend): completions by peers' st.async bytes wake a suspended try_wait late. + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(red) + : "memory"); + } + M3FA_PROF_T(27, tid - lane); + float const* mb = reinterpret_cast(smem + SX::kMailbox); + int const row = (u / kQkvSplits) * 128 + r; + // Lamport hand-off: this call's rows go to buffer (calls & 1) with relaxed stores, the consumers poll the + // values themselves (0xFFFF, a NaN the rounding of a finite sum never produces, marks "not yet"); the other + // buffer, last read by the previous call (complete: this grid passed its PDL wait), is reset to the marker. + uint16_t* cur = reinterpret_cast(p.qkv_out) + (calls & 1u) * kMaxTok * kQkvN; + uint16_t* nxt = reinterpret_cast(p.qkv_out) + ((calls & 1u) ^ 1u) * kMaxTok * kQkvN; +#pragma unroll + float qsum[kMaxTok]; +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + float4 part[kQkvSplits]; +#pragma unroll + for (int q = 0; q < kQkvSplits; ++q) + { + part[q] = *reinterpret_cast(mb + (q * kQkvOwnRows + local) * kQkvMbStride + 4 * g); + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float sum = 0.0F; +#pragma unroll + for (int q = 0; q < kQkvSplits; ++q) + { + float const pj = j == 0 ? part[q].x : j == 1 ? part[q].y : j == 2 ? part[q].z : part[q].w; + sum = sum + (q == s ? val[4 * g + j] : pj); + } + qsum[4 * g + j] = sum; + } + } +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + float const sum = qsum[t]; + if (t < p.num_tokens) + { + __nv_bfloat16 const b = __float2bfloat16_rn(sum); + uint16_t bits = *reinterpret_cast(&b); + // Never the "not yet" marker: a 0xFFFF NaN sum is stored as the NaN 0x7FFF. Only a NaN projection row + // changes (its K/V cache bytes then differ from v52 + p6 in the NaN's encoding); finite rows are exact. + bits = bits == 0xFFFFu ? static_cast(0x7FFFu) : bits; +#ifndef M3FA_QKV_REL +#define M3FA_QKV_REL 0 // 1: the row store is a release (tried: whether it becomes visible to the pollers sooner) +#endif + if (M3FA_QKV_REL) + { + asm volatile( + "st.release.gpu.global.u16 [%0], %1;\n" ::"l"(cur + static_cast(t) * kQkvN + row), + "h"(bits) + : "memory"); + } + else + { + asm volatile( + "st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(cur + static_cast(t) * kQkvN + row), + "h"(bits) + : "memory"); + } + } + asm volatile("st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(nxt + static_cast(t) * kQkvN + row), + "h"(static_cast(0xFFFFu)) + : "memory"); + } + M3FA_PROF_T(28, tid - lane); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + M3FA_PROF_T(25, 512); + if (tid == 512) + { + M3FA_CLK(46); + } +} + +// ---------------- Megakernel o_proj (v52 m3_mx_gemv bit for bit: split 2, contiguous k-tiles, nacc 1) ---------------- +// The unit's weight tiles and scale atoms into the ring (warp 0, lane 0). From L2 when the CTA prefetched them there. +template +__device__ __forceinline__ void op_issue(Params const& p, uint8_t* smem, int ou) +{ + int const rt = ou / kOpSplits; + int const kt0 = (ou % kOpSplits) * kOpTiles; + uint32_t const ob = smem_u32(smem + SX::kOBar); + for (int i = 0; i < kOpTiles; ++i) + { + mbar_expect_tx(ob + 8 * i, kQkvWTile + kQkvSfAtom); + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap_o, (kt0 + i) * 128, rt * 128, ob + 8 * i); + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf_o + (static_cast(rt) * kOpKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, ob + 8 * i); + } +} + +// L2 prefetch of the unit's weight tiles and scale atoms (one lane; nothing waits for it). +// The QKV units' share of the window-1 L2 prefetch (the next consumer's weights), issued once their own QKV weight +// stream has landed (acc_done): the attention core reads little from HBM. +__device__ __forceinline__ void window_prefetch_l2(Params const& p, int u) +{ + constexpr long long kChunk = 32768; + long long const n0 = (p.pf_b0 + kChunk - 1) / kChunk, n1 = (p.pf_b1 + kChunk - 1) / kChunk; + uint64_t pol = 0; + if (p.pf_hint) + { + asm volatile("createpolicy.fractional.L2::evict_last.b64 %0, 1.0;\n" : "=l"(pol)); + } + for (long long c = u; c < n0 + n1; c += kQkvUnits) + { + long long const a = c < n0 ? p.pf_a0 + c * kChunk : p.pf_a1 + (c - n0) * kChunk; + long long const rem = c < n0 ? p.pf_b0 - c * kChunk : p.pf_b1 - (c - n0) * kChunk; + unsigned int const n = static_cast(rem < kChunk ? rem : kChunk) & ~15u; + if (n == 0) + { + continue; + } + if (p.pf_hint) + { + asm volatile("cp.async.bulk.prefetch.L2.global.L2::cache_hint [%0], %1, %2;\n" ::"l"(a), "r"(n), "l"(pol) + : "memory"); + } + else + { + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;\n" ::"l"(a), "r"(n) : "memory"); + } + } +} + +__device__ __forceinline__ void op_prefetch_l2(Params const& p, int ou) +{ + int const rt = ou / kOpSplits; + int const kt0 = (ou % kOpSplits) * kOpTiles; + for (int i = 0; i < kOpTiles; ++i) + { + asm volatile("cp.async.bulk.prefetch.tensor.2d.L2.global [%0, {%1, %2}];\n" ::"l"(&p.wmap_o), + "r"((kt0 + i) * 128), "r"(rt * 128) + : "memory"); + } + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;\n" ::"l"( + p.wsf_o + (static_cast(rt) * kOpKTiles + kt0) * kQkvSfAtom), + "r"(kOpTiles * kQkvSfAtom) + : "memory"); +} + +// The weight tiles into TMEM as they land (warp 1, elected lane, operands warp-uniform). +template +__device__ __forceinline__ void op_stream(uint8_t* smem, uint32_t tm) +{ + uint32_t const ob = smem_u32(smem + SX::kOBar); + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint32_t const sf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kOpTiles; ++i) + { + mbar_wait_cta(ob + 8 * i, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kOA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + tmem_cp_sf(tm + kOSfa + i * 4, sf0 + i * kQkvSfAtom); + } + umma_commit(ob + 8 * kOpTiles); // a_ready +} + +// The attention's MXFP8 rows of this unit's k-tiles (the heads [kt0, kt0 + 8)), Lamport-polled value by value: 0xFF +// is NaN in both E4M3 and UE8M0, which the merge never writes (see the merge). Every thread of the CTA, one item each. +template +__device__ __forceinline__ void op_act(Params const& p, uint8_t* smem, int ou, int tid, unsigned int calls) +{ + int const kt0 = (ou % kOpSplits) * kOpTiles; + uint8_t const* q8 = p.q8 + (calls & 1u) * kMaxTok * kOpK; + uint8_t const* sf = p.sf + (calls & 1u) * kOpSfBytes; + constexpr int kXItems = kOpTiles * 8 * kMaxTok; + constexpr int kSfItems = kOpTiles * (kQkvSfAtom / 16); + for (int q = tid; q < kXItems + kSfItems; q += kThreads) + { + if (q < kXItems) + { + int const t = q % kMaxTok; + int const c = (q / kMaxTok) % 8; + int const i = q / (kMaxTok * 8); + if (t < p.num_tokens) + { + uint4 v; + bool pending; + do + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(q8 + static_cast(t) * kOpK + (kt0 + i) * 128 + c * 16) + : "memory"); + uint32_t const w[4] = {v.x, v.y, v.z, v.w}; + pending = false; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + // a byte equal to 0xFF: (x & 0x7F7F7F7F) + 0x01010101 sets bit 7 of exactly the 0x7F/0xFF bytes + uint32_t const hit = (((w[j] & 0x7F7F7F7Fu) + 0x01010101u) & w[j] & 0x80808080u); + pending |= hit != 0u; + } + } while (pending); + *reinterpret_cast(smem + SX::kOXt + i * 1024 + t * 128 + ((c ^ t) * 16)) = v; + } + } + else if (q < kXItems + kSfItems) + { + int const j = q - kXItems; + int const i = j / (kQkvSfAtom / 16); + int const c = j % (kQkvSfAtom / 16); + uint8_t const* src = sf + (kt0 + i) * kQkvSfAtom + c * 16; + uint4 v; + if (c < p.num_tokens) + { + // Row c of the atom: its 4 block scales in bytes [16 c, 16 c + 4). + do + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(src) + : "memory"); + } while ((((v.x & 0x7F7F7F7Fu) + 0x01010101u) & v.x & 0x80808080u) != 0u); + } + else + { + v = *reinterpret_cast(src); + } + *reinterpret_cast(smem + SX::kOXsf + i * kQkvSfAtom + c * 16) = v; + } + } +} + +// The MMAs (warp 1, elected lane): v52's order (per k-tile: B scale copy, 4 MMAs; one accumulator). +template +__device__ __forceinline__ void op_mma(uint8_t* smem, uint32_t tm) +{ + uint32_t const ob = smem_u32(smem + SX::kOBar); + mbar_wait_cta(ob + 8 * kOpTiles, 0); // a_ready + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + SX::kOXt)); + uint32_t const xsf0 = smem_u32(smem + SX::kOXsf); +#pragma unroll + for (int i = 0; i < kOpTiles; ++i) + { + uint32_t const sfb = tm + kOSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + umma_mx(tm + kOAcc, tm + kOA + j * 8, b0 + ((i * 1024 + kb * 32) >> 4), qkv_idesc(kb), tm + kOSfa + i * 4, + sfb, j > 0); + } + } + umma_commit(ob + 8 * (kOpTiles + 1)); // acc_done +} + +// Warps 16..19: the pair's split-K reduce (owner rows [64 s, 64 s + 64)) and the bf16 rows of y. +template +__device__ __forceinline__ void op_epilogue(Params const& p, uint8_t* smem, uint32_t tm, int tid, int ou) +{ + int const warp = tid >> 5; + int const lane = tid & 31; + uint32_t const ob = smem_u32(smem + SX::kOBar); + // The store's kernel parameters, fetched before the wait: a first touch of their constant-bank line costs a miss. + __nv_bfloat16* y_out = p.y_out; + int num_tokens = p.num_tokens; + int ar_mode = p.ar_mode; + asm volatile("" : "+l"(y_out), "+r"(num_tokens), "+r"(ar_mode)); +#ifndef M3FA_OP_PARPRE +#define M3FA_OP_PARPRE 1 // read the all-reduce slot parities before the accumulator wait (a global load off the tail) +#endif + // The parity flags are stable here: the all-reduce that flips them last completed before this grid's PDL wait, + // and the next one reads them only after this kernel's pushes. + int par_pre[kMaxTok]; + { + int const row_pre = (ou / kOpSplits) * 128 + ((tid >> 5) & 3) * 32 + (tid & 31); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + par_pre[t] = 0; + if (M3FA_OP_PARPRE && ar_mode != 0 && t < num_tokens) + { + asm volatile("ld.volatile.global.s32 %0, [%1];\n" + : "=r"(par_pre[t]) + : "l"(p.ar_flags + t * 8 + row_pre / (kOpN / 8)) + : "memory"); + } + } + } + mbar_wait_epi(ob + 8 * (kOpTiles + 1), 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF_T(51, kComputeThreads); + int const r = (warp & 3) * 32 + lane; + float v[16]; + tmem_ld16(tm + (((warp & 3) * 32) << 16) + kOAcc, v); + int const s = ou % kOpSplits; + int const owner = r / kOpOwnRows; + int const local = r % kOpOwnRows; + uint32_t const red = ob + 8 * (kOpTiles + 2); + uint32_t cta_rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + if (owner != s) + { + uint32_t const peer = cta_rank + owner - s; // o_proj units count up with the CTA index + uint32_t const dst + = mapa_shared(smem_u32(smem + SX::kOMailbox) + ((s * kOpOwnRows + local) * kQkvMbStride) * 4, peer); + uint32_t const pbar = mapa_shared(red, peer); +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];\n" ::"r"( + dst + 16 * g), + "r"(__float_as_uint(v[4 * g])), "r"(__float_as_uint(v[4 * g + 1])), "r"(__float_as_uint(v[4 * g + 2])), + "r"(__float_as_uint(v[4 * g + 3])), "r"(pbar) + : "memory"); + } + } + else + { + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(red) + : "memory"); + } + M3FA_PROF_T(52, tid - lane); + float const* mb = reinterpret_cast(smem + SX::kOMailbox); + int const row = (ou / kOpSplits) * 128 + r; + float sums[kMaxTok]; +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + float4 part[kOpSplits]; +#pragma unroll + for (int q = 0; q < kOpSplits; ++q) + { + part[q] = *reinterpret_cast(mb + (q * kOpOwnRows + local) * kQkvMbStride + 4 * g); + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float sum = 0.0F; +#pragma unroll + for (int q = 0; q < kOpSplits; ++q) + { + float const pj = j == 0 ? part[q].x : j == 1 ? part[q].y : j == 2 ? part[q].z : part[q].w; + sum = sum + (q == s ? v[4 * g + j] : pj); + } + sums[4 * g + j] = sum; + } + } + M3FA_PROF_T(54, tid - lane); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + float const sum = sums[t]; +#ifndef M3FA_EXP +#define M3FA_EXP 0 +#endif + if (t < num_tokens) + { + if (M3FA_EXP == 1) + { + asm volatile( + "st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(y_out + static_cast(t) * kOpN + row), + "h"(__bfloat16_as_ushort(__float2bfloat16_rn(sum))) + : "memory"); + } + else + { + y_out[static_cast(t) * kOpN + row] = __float2bfloat16_rn(sum); + } + } + if (ar_mode != 0) + { + // m3_mx_gemv's push into the all-reduce exchange: lane pairs hold row pairs; the even lane writes the + // bf16x2 word into this rank's slot of the (token, column slice) parity the AR's flags select. + float const hi = __shfl_xor_sync(0xffffffffu, sum, 1); + if ((lane & 1) == 0 && t < num_tokens) + { + int par; + if (M3FA_OP_PARPRE) + { + par = par_pre[t]; + } + else + { + asm volatile("ld.volatile.global.s32 %0, [%1];\n" + : "=r"(par) + : "l"(p.ar_flags + t * 8 + row / (kOpN / 8)) + : "memory"); + } + int64_t const off = static_cast(par & 1) * p.ar_maxt * p.ar_world * (kOpN / 2) + + (static_cast(t) * p.ar_world + p.ar_rank) * (kOpN / 2) + row / 2; + __nv_bfloat162 w2 = __floats2bfloat162_rn(sum, hi); + uint32_t word = *reinterpret_cast(&w2); + // -0.0 halves to +0.0: 0x80000000 marks an empty exchange word (as m3_mx_gemv's _bf16x2_word) + word = (word & 0xFFFFu) == 0x8000u ? (word & 0xFFFF0000u) : word; + word = (word >> 16) == 0x8000u ? (word & 0x0000FFFFu) : word; + asm volatile("st.global.b32 [%0], %1;\n" ::"l"(p.ar_ws + off), "r"(word) : "memory"); + } + } + } + M3FA_PROF_T(53, tid - lane); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); +} + +// Packed-QKV element group [token][col * 128 + 4 * lane, +4): the split partials, polled until every one is stored, +// summed in split order and rounded to BF16. +__device__ __forceinline__ void qkv_load(Params const& p, int par, int tok, int col, int lane, float (&e)[4]) +{ + float const* base = p.part + par * kQkvPartStride + static_cast(tok) * kQkvN + col * kHeadDim + lane * 4; + uint4 pr[kQkvSplits]; + long long spins = 0; + bool ready; + do + { + ready = true; +#pragma unroll + for (int s = 0; s < kQkvSplits; ++s) + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(pr[s].x), "=r"(pr[s].y), "=r"(pr[s].z), "=r"(pr[s].w) + : "l"(base + static_cast(s) * kMaxTok * kQkvN) + : "memory"); + ready &= pr[s].x != kQkvSentinel && pr[s].y != kQkvSentinel && pr[s].z != kQkvSentinel + && pr[s].w != kQkvSentinel; + } + if (++spins > (1LL << 26)) + { + __trap(); + } + } while (!ready); + float s0 = __uint_as_float(pr[0].x), s1 = __uint_as_float(pr[0].y), s2 = __uint_as_float(pr[0].z), + s3 = __uint_as_float(pr[0].w); +#pragma unroll + for (int s = 1; s < kQkvSplits; ++s) + { + s0 += __uint_as_float(pr[s].x); + s1 += __uint_as_float(pr[s].y); + s2 += __uint_as_float(pr[s].z); + s3 += __uint_as_float(pr[s].w); + } + e[0] = __bfloat162float(__float2bfloat16_rn(s0)); + e[1] = __bfloat162float(__float2bfloat16_rn(s1)); + e[2] = __bfloat162float(__float2bfloat16_rn(s2)); + e[3] = __bfloat162float(__float2bfloat16_rn(s3)); +} + +// Every CTA once it has read the call parity and the partials it needs: the last one flips the parity. +__device__ __forceinline__ void qkv_done(Params const& p) +{ + if (atomicAdd(p.counters + 1, 1) == kQkvGrid - 1) + { + atomicExch(p.counters + 1, 0); + atomicXor(p.counters, 1); + } +} + +} // namespace + +namespace +{ + +// FQKV: the kernel also computes the QKV(+index) projection from the layer input (see qkv_mma) and runs on a fixed +// kQkvGrid-CTA grid; the first attn_clusters clusters then run the attention on the projection's output. +#ifndef M3FA_C1_OPISSUE +#define M3FA_C1_OPISSUE \ + 4 // c1: o_proj weight TMA issued 0: at entry (before the PDL wait), 1: right after the PDL + // wait, 2: once the attention's K/V reads landed, 3: after the merge, 4: before the push, 5: after it +#endif +#ifndef M3FA_C1_TRIGGER +#define M3FA_C1_TRIGGER 3 // c1: dependents released 2: at the attention's usual point, 3: after the o_proj pushes +#endif +template +__global__ void __launch_bounds__(kThreads, 1) m3_fused_decode_attn_kernel(__grid_constant__ Params const p) +{ +#if MINIMAX_M3_ATTN_DEVICE_ENABLED + static_assert(!OPJ || (!FQKV && MXFP8), "c1: packed input, MXFP8 hand-off"); + using S = std::conditional_t, SmemX>; + constexpr bool kAttnTrigger = !FQKV && !(OPJ && M3FA_C1_TRIGGER == 3); + constexpr int kRows = S::kRows; + constexpr int kTok = S::kTok; + constexpr int kCs = kBlocks / BPC; // cluster size + constexpr int kGroups = QLEN / QG; // clusters per request + constexpr uint32_t kPushBytes = (kHeadDim + 1) * 4; // one partial row: 128 outputs and the LSE + constexpr int kTasks = kHq * QG + 3 * QLEN; // producer tasks: 16 Q heads x QG, [K, V, index-K] x QLEN + constexpr int kMaxT = (kTasks + kWarps - 1) / kWarps; + constexpr int kChunks = BPC * kPage * 4; // 16-byte packed chunks per K (or V) tile set + constexpr int kKPer = (kChunks + kComputeThreads - 1) / kComputeThreads; // K chunks per thread [0, 512) + constexpr bool kRegCore = RegCore::kOn; // warp-MMA attention core (see RegCore) + // V tiles (raw copy before the PDL wait, dequantized during QK and the softmax): on the packed path every warp + // outside the softmax groups takes part (at M=1, warps 4..19), on the fused-projection path the producer warps. +#ifndef M3FA_RS2 +#define M3FA_RS2 1 // megakernel at one 16-row chunk: the softmax and push rows split over two warp groups (8 each) +#endif + // Split softmax: the megakernel and multi-row (spec verify) steps; at one-row decode on the packed path it + // measured slower (it takes four warps from the V dequantization there). + constexpr bool kRs2 = M3FA_RS2 && BPC * QG == 1 && !RegCore::kOn && (FQKV || QLEN > 1); +#ifndef M3FA_MEGA_WARPMMA +#define M3FA_MEGA_WARPMMA 1 +#endif + // Warp-issued MMAs and the early cluster barrier: measured gains in the megakernel and on multi-row / paired-block + // steps, but not at one-row decode on 16-CTA clusters (the packed path's (1, 1, 1)), where they stay off. + constexpr bool kTuned = FQKV || QLEN > 1 || BPC > 1; + constexpr bool kWarpMma = M3FA_MEGA_WARPMMA && kTuned; + constexpr bool kEarlyClw = M3FA_MEGA_EARLYCLW && kTuned; + constexpr int kSoftmaxGroups = kRs2 ? 2 : (BPC * QG) < 4 ? (BPC * QG) : 4; +#ifndef M3FA_MEGA_VWIDE +#define M3FA_MEGA_VWIDE \ + 1 // megakernel attention CTAs (no projection unit): every warp outside the softmax copies and + // dequantizes V, as on the packed path (not only the four producer warps) +#endif + constexpr int kVStart = (kRegCore || (FQKV && !M3FA_MEGA_VWIDE)) ? kComputeThreads : 128 * kSoftmaxGroups; + constexpr int kVThreads = kThreads - kVStart; // V chunks: the warps outside the softmax groups + constexpr int kVPer = (kChunks + kVThreads - 1) / kVThreads; + constexpr int kPairs = BPC * QG; // (block, 16-row chunk) pairs of the softmax and push + constexpr int kPairsPerGroup = (kPairs + 3) / 4; // per group of four compute warps + // Named barrier 6: the softmax warps (P written) and the producer warps (V written) before the PV issue. + constexpr int kPvBarThreads = 128 * kSoftmaxGroups + kVThreads; + extern __shared__ __align__(128) uint8_t smem[]; + float* row_scale = reinterpret_cast(smem + S::kRowScale); + float* row_lse = reinterpret_cast(smem + S::kRowLse); + float* red_m = reinterpret_cast(smem + S::kRedM); + float* red_l = reinterpret_cast(smem + S::kRedL); + float* recv = reinterpret_cast(smem + S::kRecv); + uint32_t const mbar = smem_u32(smem + S::kMbar); + uint32_t const mma_bar = mbar + 8; + uint32_t* tmem_slot = reinterpret_cast(smem + S::kMbar + 16); + + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const cr = blockIdx.x % kCs; // cluster rank; owns blocks [cr * BPC, cr * BPC + BPC) + int const ci = blockIdx.x / kCs; + int const req = ci / kGroups; + int const tfirst = (ci % kGroups) * QG; // this cluster's query tokens [tfirst, tfirst + QG) of the request + int const tok0 = req * QLEN; + uint32_t const gbar = smem_u32(smem + S::kGBar); // FQKV: weight ring, A ready, accumulators done + int const qkv_u = FQKV ? qkv_unit() : -1; // FQKV: this CTA's projection unit, or -1 + + // ---------------- Prologue: nothing here depends on the producer kernel ---------------- +#if M3FA_TRIGGER_AT == 0 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (!FQKV) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + M3FA_PROF(0); +#ifdef M3FA_PROFILE + if (p.prof != nullptr && tid == 0) + { + uint32_t smid; + asm volatile("mov.u32 %0, %%smid;" : "=r"(smid)); + p.prof[blockIdx.x * 64 + 40] = smid; + } +#endif + if constexpr (FQKV) + { + // The projection weights stream into TMEM while the producer (the all-reduce) still runs. + if (warp < 2) + { + qkv_setup(p, smem, gbar, tmem_slot, warp, lane, qkv_u); + } + if (ci >= p.attn_clusters) + { + // A projection-only CTA. Every reduce barrier of the cluster is initialized before any peer pushes. + asm volatile("barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;\n" :: + : "memory"); + int const op_u = static_cast(blockIdx.x) - kOpFirstCta; // this CTA's o_proj unit (>= 0 here) +#ifndef M3FA_OP_EARLY +#define M3FA_OP_EARLY 0 // 1: o_proj-only CTAs stream their weights from launch (0: after the PDL wait) +#endif + if (M3FA_OP_EARLY && qkv_u < 0 && warp == 0 && lane == 0) + { + op_issue(p, smem, op_u); // an o_proj-only CTA: its weights stream from launch on + } + if (M3FA_OP_L2PF && qkv_u >= 0 && warp == 0 && lane == 0) + { + op_prefetch_l2(p, op_u); // into L2 during the all-reduce, behind this CTA's QKV ring + } + qkv_stream(p, smem, gbar, tmem_slot, warp, lane, qkv_u); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (M3FA_MEGA_TRIGGER == 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + // This CTA's call counter: read after the PDL wait (the previous call's CTA increments it at its end), and + // early, so its latency stays off the epilogue's critical path. + unsigned int const calls = p.flags[32 + blockIdx.x]; + M3FA_PROF(1); + if (!M3FA_OP_EARLY && qkv_u < 0 && warp == 0 && lane == 0) + { + op_issue(p, smem, op_u); + } + if (qkv_u >= 0) + { + int const par = 0; + qkv_mma(p, smem, gbar, tmem_slot, tid, qkv_u); + if (warp >= kComputeWarps) + { + qkv_epilogue(p, smem, gbar, tmem_slot, tid, qkv_u, par, calls); + } + if (warp == 0 && lane == 0) + { + // The ring (and the activation over it) is free once the QKV MMAs completed. + mbar_wait_cta(gbar + 16 * kQkvRing + 8, 0); + op_issue(p, smem, op_u); + if (p.pf_b0 > 0 || p.pf_b1 > 0) + { + window_prefetch_l2(p, qkv_u); + } + } + } + __syncthreads(); // the QKV accumulator reads are done: TMEM columns [0, 512) are the o_proj's + M3FA_PROF(2); + uint32_t const tm = *tmem_slot; + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if (elect_one()) + { + op_stream(smem, tm); + } + __syncwarp(); + } + op_act(p, smem, op_u, tid, calls); + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + M3FA_PROF_T(36, 0); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Late release: the dependents launch once every unit has its attention input (the rel_v4 rule). + if (M3FA_MEGA_TRIGGER < 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + __syncthreads(); + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if (elect_one()) + { + op_mma(smem, tm); + M3FA_PROF_T(50, threadIdx.x); + } + __syncwarp(); + } + if (warp >= kComputeWarps) + { + op_epilogue(p, smem, tm, tid, op_u); + } + M3FA_PROF_T(37, kComputeThreads); + __syncthreads(); + M3FA_PROF_T(55, 0); + if (M3FA_MEGA_TRIGGER == 3) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // after this CTA's y rows / pushes + } + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" ::"r"(*tmem_slot), "r"(kTmemAll) + : "memory"); + M3FA_PROF_T(56, 32); + } + if (tid == 0) + { + p.flags[32 + blockIdx.x] = calls + 1u; // a plain store: no load on the exit path + } + M3FA_PROF(39); +#ifndef M3FA_FINAL_CLBAR +#define M3FA_FINAL_CLBAR 0 +#endif + if (M3FA_FINAL_CLBAR) + { + // No peer may exit while its partials can still be pushed into it (DSMEM lifetime). + asm volatile("barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;\n" :: + : "memory"); + } + M3FA_PROF(38); + return; + } + } +#ifndef M3FA_MEGA_ATTN_WAIT_FIRST +#define M3FA_MEGA_ATTN_WAIT_FIRST 0 // 1: the attention CTAs read nothing (metadata, caches) before their PDL wait +#endif +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (FQKV && M3FA_MEGA_ATTN_WAIT_FIRST) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } +#endif + if (!FQKV && warp == kWarps - 1) // a producer warp: it has nothing else to do before the PDL wait + { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n" ::"r"(smem_u32(tmem_slot)), + "r"(S::kTmemCols) + : "memory"); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n" ::: "memory"); + } + int const seq_len = p.seq_lens[req]; + int page[BPC]; +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + int const b = cr * BPC + h; + page[h] = b < p.bt_cols ? p.block_table[static_cast(req) * p.bt_stride + b] : 0; + } + // New-token cache slots, for the tokens up to this cluster's last one. + int slot[QLEN]; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + slot[t] = t < tfirst + QG ? p.out_cache_loc[tok0 + t] : -1; + } + // This warp's producer tasks and, for every lane that applies RoPE, the task token's cos/sin. + float cs[kMaxT][4], sn[kMaxT][4]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + int const t = i < kHq * QG ? tfirst + i / kHq : (i - kHq * QG) / 3; + if (i < kTasks && lane < 16 && t < QLEN) + { + int64_t const row = static_cast(p.position_ids[tok0 + t]) * kRotDim; +#pragma unroll + for (int e = 0; e < 4; ++e) + { + int const coefficient = (lane * 4 + e) & (kRotDim / 2 - 1); + cs[k][e] = p.rope[row + coefficient]; + sn[k][e] = p.rope[row + kRotDim / 2 + coefficient]; + } + } + else + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + cs[k][e] = 1.0F; + sn[k][e] = 0.0F; + } + } + } + float qw[4], kw[4], ikw[4]; + load_w4(qw, p.q_w, lane); + load_w4(kw, p.k_w, lane); + load_w4(ikw, p.ik_w, lane); + float const kqs = p.kv_quant_scale[1]; + float const vqs = p.kv_quant_scale[2]; + float const sl2e = (p.sm_scale * 1.4426950409f) * p.k_gs[0]; + float const vgs = p.v_gs[0]; + // Longest attended length among this cluster's tokens, and each block's rows within it. + int const kv_max = max(seq_len - QLEN + tfirst + QG, 0); + int nvh[BPC]; +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + nvh[h] = min(max(kv_max - (cr * BPC + h) * kPage, 0), kPage); + } + bool const any_valid = nvh[0] > 0; + if (tid == 32) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(mbar), "r"(1)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(mma_bar), "r"(1)); + asm volatile( + "mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(mbar), "r"(kBlocks * QG * BPC * kPushBytes) + : "memory"); + if constexpr (OPJ) + { + uint32_t const ob = smem_u32(smem + S::kOBar); + for (int i = 0; i < kOpTiles + 3; ++i) + { + mbar_init(ob + 8 * i, 1); + } + mbar_expect_tx(ob + 8 * (kOpTiles + 2), (kOpSplits - 1) * kOpOwnRows * kMaxTok * 4); + } + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + if constexpr (OPJ) + { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if (M3FA_C1_OPISSUE == 0) + { + op_issue(p, smem, static_cast(blockIdx.x)); + } + } + } + // This thread's packed K / V chunks, held in registers until they are dequantized after the PDL wait. + RawChunk kraw[kKPer]; +#pragma unroll + for (int k = 0; k < kKPer; ++k) + { + int const q = tid + kComputeThreads * k; + kraw[k] = RawChunk{make_uint4(0u, 0u, 0u, 0u), 0u}; + int const h = q / (kPage * 4); + if (q < kChunks && tid < kComputeThreads && nvh[h] > 0) + { + kraw[k] = load_chunk(p.kv_data + static_cast(page[h]) * p.d_page, + p.kv_scale + static_cast(page[h]) * p.s_page, chunk_row(q % (kPage * 4)), + chunk_col(q % (kPage * 4)), false); + } + } + // V goes through cp.async into shared memory: an asynchronous copy does not hold the producer warps up the way + // an outstanding register load does, and they need V only after QK has been issued. + if (tid >= kVStart) + { +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + if (nvh[h] > 0) + { + uint8_t* dst = smem + S::kRawV + h * S::kRawVBlock; + uint8_t const* data = p.kv_data + static_cast(page[h]) * p.d_page + p.d_plane; + uint8_t const* scales = p.kv_scale + static_cast(page[h]) * p.s_page + p.s_plane; + for (int c = tid - kVStart; c < S::kRawVBlock / 16; c += kVThreads) + { + cp_async16(dst + c * 16, c < kPage * 4 ? data + c * 16 : scales + (c - kPage * 4) * 16); + } + } + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + if (warp == kWarps - 1) + { + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + // Every CTA's mbarrier is initialized before any peer pushes into it; the matching wait sits right before the + // push, so this barrier's latency hides behind the attention (with aliased partials it moves there). + if constexpr (!S::kRecvAlias) + { + asm volatile("barrier.cluster.arrive.relaxed.aligned;\n" ::: "memory"); + } + int kv_len[QG]; +#pragma unroll + for (int tl = 0; tl < QG; ++tl) + { + kv_len[tl] = max(seq_len - QLEN + tfirst + tl + 1, 0); + } + int const cta_tok0 = cr * kTok; // request position of this CTA's first block row + // New token t's row in this CTA's tiles (block * 128 + row), or -1: the CTA owns the block holding the token's + // position, and the token is attended by this cluster (t <= its last token). Every thread derives the same. + int new_row[QLEN]; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + int const pos = seq_len - QLEN + t; + int const blk = min(pos >> 7, kBlocks - 1); + int const spage = slot[t] / kPage; + bool const mine = t < tfirst + QG && pos >= 0 && blk / BPC == cr && slot[t] >= 0 && spage < p.num_pages; + new_row[t] = mine && spage == page[blk % BPC] ? (blk % BPC) * kPage + (slot[t] & (kPage - 1)) : -1; + } + auto is_new_row = [&](int r) + { + bool hit = false; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + hit |= new_row[t] == r; + } + return hit; + }; + // Cache writes of the new K/V/index-K are issued after this CTA has pushed its partials: only the fp16 tile rows + // are on the attention's path. kind 0: none, 1: NVFP4 data + scale, 2: FP8 index-K. + int wr_kind[kMaxT]; + uint8_t* wr_ptr[kMaxT]; + uint8_t* wr_sptr[kMaxT]; + uint32_t wr_val[kMaxT]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + wr_kind[k] = 0; + wr_ptr[k] = nullptr; + wr_sptr[k] = nullptr; + wr_val[k] = 0u; + } + uint32_t tmem = 0; + auto dequant_k = [&]() + { + if (tid < kComputeThreads) + { +#pragma unroll + for (int k = 0; k < kKPer; ++k) + { + int const q = tid + kComputeThreads * k; + if (q < kChunks) + { + int const h = q / (kPage * 4); + int const row = chunk_row(q % (kPage * 4)); + if (!is_new_row(h * kPage + row)) + { + store_chunk_fp16( + smem + S::kK16 + h * S::kTileBytes, row, chunk_col(q % (kPage * 4)), kraw[k], row < nvh[h]); + } + } + } + } + }; + // The producer warps' V tiles (raw tiles from cp.async, dequantized to fp16), then their PV-barrier arrival. + auto dequant_v = [&]() + { +#ifdef M3FA_DBG_NO_VDQ + if constexpr (!kRegCore) + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + return; +#endif + asm volatile("cp.async.wait_all;\n" ::: "memory"); + asm volatile("bar.sync 4, %0;\n" ::"n"(kVThreads) : "memory"); +#pragma unroll + for (int k = 0; k < kVPer; ++k) + { + int const q = tid - kVStart + kVThreads * k; + if (q < kChunks) + { + int const h = q / (kPage * 4); + int const row = chunk_row(q % (kPage * 4)); + int const cc = chunk_col(q % (kPage * 4)); + if (!is_new_row(h * kPage + row)) + { + uint8_t const* rawv = smem + S::kRawV + h * S::kRawVBlock; + uint8_t const* sc = rawv + kPage * 64 + (row >> 2) * 32 + (row & 3); + RawChunk c; + c.data = *reinterpret_cast(rawv + row * 64 + cc * 16); + c.scales + = static_cast(sc[(cc * 2) * 4]) | (static_cast(sc[(cc * 2 + 1) * 4]) << 8); + store_chunk_fp16(smem + S::kV16 + h * S::kTileBytes, row, cc, c, row < nvh[h]); + } + } + } + M3FA_FENCE_PROXY_ASYNC(); // V tiles -> tensor core reads (ordered by the barrier before the PV issue) + if constexpr (!kRegCore) + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + M3FA_PROF_T(16, kComputeThreads); + M3FA_PROF_T(17, kThreads - 1); + }; + if constexpr (FQKV) + { + qkv_stream(p, smem, gbar, tmem_slot, warp, lane, qkv_u); + } + if constexpr (!FQKV) + { + if (M3FA_V15_EARLYV && any_valid && tid >= kVStart) + { + dequant_v(); + } + if constexpr (kEarlyClw && !S::kRecvAlias) + { + asm volatile("barrier.cluster.wait.aligned;\n" ::: "memory"); + } + } + M3FA_PROF(11); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (FQKV && M3FA_MEGA_TRIGGER == 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#if M3FA_TRIGGER_AT == 1 + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +#endif +#endif + // FQKV (megakernel): this call's QKV rows sit in buffer (calls & 1), polled value by value below. + unsigned int const mega_calls = (FQKV || OPJ) ? p.flags[32 + blockIdx.x] : 0u; + if (OPJ && M3FA_C1_OPISSUE == 1 && tid == 32) + { + op_issue(p, smem, static_cast(blockIdx.x)); + } + __nv_bfloat16 const* qkv_src = FQKV ? p.packed + (mega_calls & 1u) * kMaxTok * kQkvN : p.packed; + M3FA_PROF(1); + int const par = 0; + if constexpr (FQKV) + { + if (qkv_u >= 0) + { + qkv_mma(p, smem, gbar, tmem_slot, tid, qkv_u); + } + } + + if constexpr (FQKV) + { + if (qkv_u >= 0) + { + if (warp >= kComputeWarps) + { + qkv_epilogue(p, smem, gbar, tmem_slot, tid, qkv_u, par, mega_calls); + } + else + { + // The attention regions overlay the weight ring: every weight byte has gone on into TMEM (or, with + // the MMAs reading the ring, every MMA has completed) first. + mbar_wait_cta(gbar + 16 * kQkvRing + (kQkvASmem ? 8 : 0), 0); + } + } + if (any_valid) + { + dequant_k(); + } +#ifndef M3FA_MEGA_EARLYV +#define M3FA_MEGA_EARLYV 1 // attention CTAs dequantize V here, while the QKV rows are still on their way +#endif +#ifndef M3FA_MEGA_EARLYCLW +#define M3FA_MEGA_EARLYCLW 1 // and pass the cluster barrier (peers' merge mbarriers initialized) here too +#endif + if (M3FA_MEGA_EARLYV && any_valid && tid >= kVStart) + { + dequant_v(); + } + if constexpr (kEarlyClw && !S::kRecvAlias) + { + asm volatile("barrier.cluster.wait.aligned;\n" ::: "memory"); + } + M3FA_PROF(14); + } + + if (any_valid) // CTA-uniform + { + // ---------------- Producer ---------------- + // Q rows of this cluster's tokens in every CTA. The new K/V of every token up to this cluster's last one + // whose position falls in this CTA's blocks go straight into the fp16 tiles; the CTA writes the caches, and + // the index-K, only for its own tokens. + uint2 raw[kMaxT]; + float fq[FQKV ? kMaxT : 1][4]; + bool act[kMaxT]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + raw[k] = make_uint2(0u, 0u); + act[k] = false; + if (i < kTasks) + { + int t, col; + bool active = true; + if (i >= kHq * QG) + { + int const kv = i - kHq * QG; + t = kv / 3; + int const kind = kv % 3; + col = kind == 0 ? kColK : (kind == 1 ? kColV : kColIdxK); + int const pos = seq_len - QLEN + t; + active = t < tfirst + QG && pos >= 0 && min(pos >> 7, kBlocks - 1) / BPC == cr && slot[t] >= 0 + && slot[t] / kPage < p.num_pages && (kind != 2 || t >= tfirst); + } + else + { + t = tfirst + i / kHq; + col = i % kHq; + } + act[k] = active; + if (active) + { + uint2 const* src = reinterpret_cast( + qkv_src + static_cast(tok0 + t) * p.packed_stride + col * kHeadDim + lane * 4); + if constexpr (FQKV) + { + uint2 v; + bool pending; + do + { + asm volatile("ld.relaxed.gpu.global.v2.b32 {%0, %1}, [%2];\n" + : "=r"(v.x), "=r"(v.y) + : "l"(src) + : "memory"); + pending = (v.x & 0xFFFFu) == 0xFFFFu || (v.x >> 16) == 0xFFFFu || (v.y & 0xFFFFu) == 0xFFFFu + || (v.y >> 16) == 0xFFFFu; + } while (pending); + raw[k] = v; + } + else + { + raw[k] = *src; + } + } + } + } + M3FA_PROF(15); + // While the packed-QKV loads are in flight: dequantize this thread's K chunks to fp16 (exact: e2m1 x e4m3 + // fits fp16). Rows past the longest attended length are zeroed, as the Triton producer masks its loads; the + // new tokens' rows are written by their producer tasks instead. (Building the tiles before the PDL wait was + // measured slower: the cache loads issued at entry are still in flight there, so the wait moves later.) + if constexpr (!FQKV) + { + dequant_k(); + } +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + if (i >= kTasks || !act[k]) + { + continue; + } + float e[4]; + { + float2 const a = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw[k].x)); + float2 const b = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw[k].y)); + e[0] = a.x; + e[1] = a.y; + e[2] = b.x; + e[3] = b.y; + } + if (i < kHq * QG) + { + norm_rope(e, qw, p.eps, cs[k], sn[k], lane); + __nv_fp8x2_e4m3 const lo(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const hi(make_float2(e[2], e[3])); + uint2 hv; + hv.x = e4m3x2_to_h2(lo.__x); + hv.y = e4m3x2_to_h2(hi.__x); + *reinterpret_cast(smem + S::kQ + kmaj_off(i, lane * 4)) = hv; // row i = token * 16 + head + continue; + } + int const t = (i - kHq * QG) / 3; + int const kind = (i - kHq * QG) % 3; + int const spage = slot[t] / kPage; + int const within = slot[t] & (kPage - 1); + if (kind == 2) + { + norm_rope(e, ikw, p.eps, cs[k], sn[k], lane); + __nv_fp8x2_e4m3 const lo(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const hi(make_float2(e[2], e[3])); + wr_kind[k] = 2; + wr_val[k] = static_cast(lo.__x) | (static_cast(hi.__x) << 16); + wr_ptr[k] = p.idx_k + static_cast(spage) * p.i_page + static_cast(within) * p.i_tok + + lane * 4; + continue; + } + bool const isV = kind == 1; + if (!isV) + { + norm_rope(e, kw, p.eps, cs[k], sn[k], lane); + } + uint8_t sbyte = 0; + uint16_t const fp4 = quantizeMinimaxM3Fp4x4(e, isV ? vqs : kqs, &sbyte); + if (t >= tfirst) + { + int const scale_col = lane / 4; + int const scale_within = isV ? (within / 4) * 32 + scale_col * 4 + within % 4 : within * 8 + scale_col; + wr_kind[k] = 1; + wr_val[k] = static_cast(fp4) | (static_cast(sbyte) << 16); + wr_ptr[k] = p.kv_data + static_cast(spage) * p.d_page + (isV ? p.d_plane : 0) + within * 64 + + lane * 2; + wr_sptr[k] = p.kv_scale + static_cast(spage) * p.s_page + (isV ? p.s_plane : 0) + scale_within; + } + if (new_row[t] >= 0) + { + // The attention reads this page: its fp16 tile row is the freshly quantized row, dequantized like + // the rest of the tile. + int const h = new_row[t] / kPage; + int const row = new_row[t] % kPage; + uint32_t const sc2 = e4m3_to_h2_splat(sbyte); + uint2 hv; + hv.x = hmul2_u(e2m1x2_to_h2(fp4 & 0xFFu), sc2); + hv.y = hmul2_u(e2m1x2_to_h2(fp4 >> 8), sc2); + *reinterpret_cast( + smem + (isV ? S::kV16 : S::kK16) + h * S::kTileBytes + kmaj_off(row, lane * 4)) + = hv; + } + } + // K tiles, Q rows and patched rows -> tensor core reads. Only threads that wrote fence: the fence also + // waits for the thread's outstanding loads, and the producer warps still have their V chunks in flight. + bool wrote = tid < kComputeThreads; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + wrote |= act[k]; + } + if (wrote) + { + M3FA_FENCE_PROXY_ASYNC(); + } + M3FA_PROF(13); + __syncthreads(); +#if M3FA_TRIGGER_AT == 2 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Every K/V byte this CTA reads has arrived (the V chunks sit in registers): release the dependents now so + // their weight streams do not queue ahead of the KV reads (the megakernel releases late, from its o_proj). + if constexpr (kAttnTrigger || (FQKV && M3FA_MEGA_TRIGGER == 1)) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + if (OPJ && M3FA_C1_OPISSUE == 2 && tid == 32) + { + op_issue(p, smem, static_cast(blockIdx.x)); + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + tmem = *tmem_slot; + M3FA_PROF(2); + // ---------------- S^T = K Q^T per block on the tensor core: TMEM lane = token, column = query row ---------- +#ifndef M3FA_MEGA_WARPMMA +#define M3FA_MEGA_WARPMMA 1 // attention MMAs issued by warp 0's elected lane with warp-uniform operands (not tid 0) +#endif + if (!kRegCore && (kWarpMma ? warp == 0 : tid == 0)) + { + // Operands computed by the whole warp stay in uniform registers; one elected lane issues. + uint32_t const tm_u = tmem; + bool const issue = kWarpMma ? elect_one() : true; + if (issue) + { +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + uint32_t const a0 = smem_u32(smem + S::kK16 + h * S::kTileBytes); + uint32_t const b0 = smem_u32(smem + S::kQ); +#pragma unroll + for (int ks = 0; ks < kHeadDim / 16; ++ks) + { + umma_f16(tm_u + h * kRows, umma_desc(a0 + ks * 256, 128, 2048), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 0, 0), ks > 0); + } + } + umma_commit(mma_bar); + } + if (kWarpMma) + { + __syncwarp(); + } + } + if (tid >= kVStart && !(FQKV && M3FA_MEGA_EARLYV) && !(!FQKV && M3FA_V15_EARLYV)) + { + // The producer warps dequantize V while QK and the softmax run. + dequant_v(); + } + if constexpr (kRegCore) + { + // ---------------- Register core: warp w of block h = w / kWpb owns 8 * kNt tokens for S and 8 * kNt + // output dims for O; m16n8k16 fp16 MMAs with fp32 accumulators, the softmax in the S fragments ---------- + constexpr int kWpb = kComputeWarps / BPC; // warps per block + constexpr int kNt = 16 / kWpb; // 8-token (S) or 8-dim (O) tiles per warp + float* redx_m = reinterpret_cast(smem + S::kRedX); + float* redx_l = redx_m + BPC * kRows * kWpb; + if (warp < kComputeWarps) + { + int const h = warp / kWpb; + int const wib = warp % kWpb; + int const g = lane >> 2; + int const q4 = lane & 3; + uint32_t const ktile = smem_u32(smem + S::kK16 + h * S::kTileBytes); + uint32_t const vtile = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const qrows = smem_u32(smem + S::kQ); + uint8_t* ptile = smem + S::kK16 + h * S::kTileBytes; // P over this block's K tile, after QK + auto block_bar = [&]() + { + if (h == 0) + { + asm volatile("bar.sync 14, %0;\n" ::"n"(kWpb * 32) : "memory"); + } + else + { + asm volatile("bar.sync 15, %0;\n" ::"n"(kWpb * 32) : "memory"); + } + }; + // S = Q K^T: rows (c, g / g + 8), tokens 8 tb + 2 q4 (+1). + float sv[QG][kNt][4]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + sv[c][nt][e] = 0.0F; + } + } + } +#pragma unroll + for (int ks = 0; ks < kHeadDim / 16; ++ks) + { + uint32_t bk[kNt][2]; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const tb = wib * kNt + nt; + ldsm_x2(ktile + tb * 2048 + (ks * 2 + ((lane >> 3) & 1)) * 128 + (lane & 7) * 16, bk[nt]); + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { + uint32_t aq[4]; + ldsm_x4( + qrows + (c * 2 + ((lane >> 3) & 1)) * 2048 + (ks * 2 + (lane >> 4)) * 128 + (lane & 7) * 16, + aq); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + mma_f16_16816(sv[c][nt], aq, bk[nt]); + } + } + } + M3FA_PROF(7); + // Row maxima (log2 domain, masked as the producer masks), quad then across the block's warps. + float mrow[QG][2]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { + float mx0 = kNegInf, mx1 = kNegInf; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + int const tok = cta_tok0 + h * kPage + (wib * kNt + nt) * 8 + 2 * q4 + (e & 1); + bool const valid = tok < kv_max && tok < kv_len[c]; + float const v = valid ? sv[c][nt][e] * sl2e : kNegInf; + sv[c][nt][e] = v; + if (e < 2) + { + mx0 = fmaxf(mx0, v); + } + else + { + mx1 = fmaxf(mx1, v); + } + } + } + mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, 1)); + mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, 2)); + mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, 1)); + mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, 2)); + if (q4 == 0) + { + redx_m[(h * kRows + c * 16 + g) * kWpb + wib] = mx0; + redx_m[(h * kRows + c * 16 + g + 8) * kWpb + wib] = mx1; + } + } + block_bar(); // every warp of the block is past QK (P may overwrite the K tile) and wrote its maxima + M3FA_PROF(8); +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int half = 0; half < 2; ++half) + { + float const* rm = redx_m + (h * kRows + c * 16 + g + 8 * half) * kWpb; + float m = kNegInf; +#pragma unroll + for (int w4 = 0; w4 < kWpb; w4 += 4) + { + float4 const v = *reinterpret_cast(rm + w4); + m = fmaxf(m, fmaxf(fmaxf(v.x, v.y), fmaxf(v.z, v.w))); + } + mrow[c][half] = m; + } + // P = exp2(s - m) as fp16 into [row][token] core matrices; the fp32 values feed the row sums. + float l0 = 0.0F, l1 = 0.0F; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + float a[4]; +#pragma unroll + for (int e = 0; e < 4; ++e) + { + float const m = mrow[c][e >> 1]; + a[e] = m == kNegInf ? 0.0F : ex2(sv[c][nt][e] - m); + } + l0 += a[0] + a[1]; + l1 += a[2] + a[3]; + int const t0 = (wib * kNt + nt) * 8 + 2 * q4; + int const r0 = c * 16 + g; + __half2 const p0 = __floats2half2_rn(a[0], a[1]); + __half2 const p1 = __floats2half2_rn(a[2], a[3]); + *reinterpret_cast<__half2*>( + ptile + (r0 >> 3) * 2048 + (t0 >> 3) * 128 + (r0 & 7) * 16 + (t0 & 7) * 2) + = p0; + *reinterpret_cast<__half2*>( + ptile + ((r0 + 8) >> 3) * 2048 + (t0 >> 3) * 128 + (r0 & 7) * 16 + (t0 & 7) * 2) + = p1; + } + l0 += __shfl_xor_sync(0xffffffffu, l0, 1); + l0 += __shfl_xor_sync(0xffffffffu, l0, 2); + l1 += __shfl_xor_sync(0xffffffffu, l1, 1); + l1 += __shfl_xor_sync(0xffffffffu, l1, 2); + if (q4 == 0) + { + redx_l[(h * kRows + c * 16 + g) * kWpb + wib] = l0; + redx_l[(h * kRows + c * 16 + g + 8) * kWpb + wib] = l1; + } + } + M3FA_PROF_T(19, 0); + } + __syncthreads(); // P, the V tiles (producer warps) and the row sums visible + M3FA_PROF(9); + if (warp < kComputeWarps) + { + int const h = warp / kWpb; + int const wib = warp % kWpb; + int const g = lane >> 2; + int const q4 = lane & 3; + uint32_t const vtile = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const ptile = smem_u32(smem + S::kK16 + h * S::kTileBytes); + float* obuf = reinterpret_cast(smem + S::kK16 + h * S::kTileBytes + S::kObufOff); + // O = P V: rows (c, g / g + 8), dims 8 dt + 2 q4 (+1). + float ov[QG][kNt][4]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + ov[c][nt][e] = 0.0F; + } + } + } +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + uint32_t bv[kNt][2]; + int const tok = ks * 16 + ((lane >> 3) & 1) * 8 + (lane & 7); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const dt = wib * kNt + nt; + ldsm_x2_trans(vtile + (tok >> 3) * 2048 + dt * 128 + (tok & 7) * 16, bv[nt]); + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { + uint32_t ap[4]; + ldsm_x4( + ptile + (c * 2 + ((lane >> 3) & 1)) * 2048 + (ks * 2 + (lane >> 4)) * 128 + (lane & 7) * 16, + ap); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + mma_f16_16816(ov[c][nt], ap, bv[nt]); + } + } + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const d = (wib * kNt + nt) * 8 + 2 * q4; + int const r0 = c * 16 + g; + *reinterpret_cast(obuf + r0 * kOPitch + d) = make_float2(ov[c][nt][0], ov[c][nt][1]); + *reinterpret_cast(obuf + (r0 + 8) * kOPitch + d) + = make_float2(ov[c][nt][2], ov[c][nt][3]); + } + } + } + else if (tid - kComputeThreads < BPC * kRows) + { + // Row statistics of every (block, row): LSE and the PV scale, as the producer derives them. + int const r = tid - kComputeThreads; + float const* rm = redx_m + r * kWpb; + float const* rl = redx_l + r * kWpb; + float m = kNegInf; + float l = 0.0F; +#pragma unroll + for (int w4 = 0; w4 < kWpb; w4 += 4) + { + float4 const vm = *reinterpret_cast(rm + w4); + float4 const vl = *reinterpret_cast(rl + w4); + m = fmaxf(m, fmaxf(fmaxf(vm.x, vm.y), fmaxf(vm.z, vm.w))); + l += (vl.x + vl.y) + (vl.z + vl.w); + } + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[r] = scale * vgs; + row_lse[r] = lse; + } + __syncthreads(); // O partials and row statistics visible to the push + } + else + { + if constexpr (kRs2) + { + // ---------------- Split softmax: warps 0..3 rows 0..7, warps 4..7 rows 8..15 (thread = block token) + // -------- + if (warp < 8) + { + int const rg = warp / 4; + int const wq = warp % 4; + int const tk = wq * 32 + lane; + auto rg_bar = [&]() + { + if (rg == 0) + { + asm volatile("bar.sync 10, 128;\n" ::: "memory"); + } + else + { + asm volatile("bar.sync 11, 128;\n" ::: "memory"); + } + }; + int const rs8 = ((lane >> 4) & 1) * 4 + ((lane >> 3) & 1) * 2 + ((lane >> 2) & 1); + mbar_wait_cta(mma_bar, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF(7); + float v[8]; + tmem_ld8(tmem + ((wq * 32) << 16) + 8 * rg, v); + bool const valid = cta_tok0 + tk < kv_max && cta_tok0 + tk < kv_len[0]; + float a[8]; +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] = valid ? v[j2] * sl2e : kNegInf; + a[j2] = v[j2]; + } + float const mw = warp_reduce_scatter8(a, lane); + if ((lane & 3) == 0) + { + red_m[(8 * rg + rs8) * 4 + wq] = mw; + } + rg_bar(); + M3FA_PROF(8); + uint8_t* pdst = smem + S::kP + (tk >> 3) * 128 + (tk & 7) * 16; + float mrow = kNegInf; + if (lane < 8) + { + float4 const mq = *reinterpret_cast(red_m + (8 * rg + lane) * 4); + mrow = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + } + float pa[8]; +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + float const m = __shfl_sync(0xffffffffu, mrow, j2); + pa[j2] = m == kNegInf ? 0.0F : ex2(v[j2] - m); + } + if (tid == 0) + { + M3FA_CLK(41); + } + uint32_t pw[4]; +#pragma unroll + for (int j2 = 0; j2 < 4; ++j2) + { + __half2 const hp = __floats2half2_rn(pa[2 * j2], pa[2 * j2 + 1]); + pw[j2] = *reinterpret_cast(&hp); + } + *reinterpret_cast(pdst + rg * 2048) = make_uint4(pw[0], pw[1], pw[2], pw[3]); + M3FA_PROF_T(18, 0); + if (tid == 0) + { + M3FA_CLK(42); + } + M3FA_FENCE_PROXY_ASYNC(); // P -> tensor core reads + if (tid == 0) + { + M3FA_CLK(43); + } + if (warp == 0) + { + asm volatile("bar.sync 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + if (tid == 0) + { + M3FA_CLK(44); + } + } + else + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + if (kWarpMma ? warp == 0 : tid == 0) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm_u = tmem; + uint32_t const a0 = smem_u32(smem + S::kV16); + uint32_t const b0 = smem_u32(smem + S::kP); + if (kWarpMma ? elect_one() : true) + { +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + umma_f16(tm_u + BPC * kRows, umma_desc(a0 + ks * 4096, 2048, 128), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 1, 1), ks > 0); + } + umma_commit(mma_bar); + } + if (kWarpMma) + { + __syncwarp(); + } + } + M3FA_PROF_T(19, 0); + float const lw = warp_reduce_scatter8(pa, lane); + if ((lane & 3) == 0) + { + red_l[(8 * rg + rs8) * 4 + wq] = lw; + } + rg_bar(); + M3FA_PROF_T(20, 0); + if (wq == 0 && lane < 8) + { + int const r = 8 * rg + lane; + float4 const mq = *reinterpret_cast(red_m + r * 4); + float4 const lq = *reinterpret_cast(red_l + r * 4); + float const m = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + float const l = (lq.x + lq.y) + (lq.z + lq.w); + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[r] = scale * vgs; + row_lse[r] = lse; + } + } + __syncthreads(); // row scales and LSEs visible (PV was issued above) + M3FA_PROF(9); + if (warp < 8) + { + mbar_wait_cta(mma_bar, 1); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + } + else + { + // ---------------- Softmax: thread = block token; the (block, 16-row chunk) pairs are spread over the + // four groups of four compute warps (each group covers the 128 TMEM lanes = block tokens) + // ---------------- + if (warp < kComputeWarps && warp / 4 < kPairs) + { + int const g = warp / 4; + int const wq = warp % 4; + int const tk = wq * 32 + lane; + auto group_bar = [&]() + { + if (g == 0) + { + asm volatile("bar.sync 10, 128;\n" ::: "memory"); + } + else if (g == 1) + { + asm volatile("bar.sync 11, 128;\n" ::: "memory"); + } + else if (g == 2) + { + asm volatile("bar.sync 12, 128;\n" ::: "memory"); + } + else + { + asm volatile("bar.sync 13, 128;\n" ::: "memory"); + } + }; + // Row r(lane) this lane holds after the reduce-scatter: bits 4..1 of the lane. + int const rs_row + = ((lane >> 4) & 1) * 8 + ((lane >> 3) & 1) * 4 + ((lane >> 2) & 1) * 2 + ((lane >> 1) & 1); + mbar_wait_cta(mma_bar, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF(7); + float v[kPairsPerGroup][16]; +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + tmem_ld16(tmem + ((wq * 32) << 16) + h * kRows + c * 16, v[k]); + bool const valid + = cta_tok0 + h * kPage + tk < kv_max && cta_tok0 + h * kPage + tk < kv_len[c]; + float a[16]; +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[k][j] = valid ? v[k][j] * sl2e : kNegInf; + a[j] = v[k][j]; + } + float const mw = warp_reduce_scatter16(a, lane); + if ((lane & 1) == 0) + { + red_m[((h * QG + c) * 16 + rs_row) * 4 + wq] = mw; + } + } + } + group_bar(); + M3FA_PROF(8); + float pa[kPairsPerGroup][16]; +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + uint8_t* pdst = smem + S::kP + h * kRows * kPage * 2 + (tk >> 3) * 128 + (tk & 7) * 16; + // Row maxima: lane j (< 16) reduces row j's four warp partials, then broadcasts it. + float mrow = kNegInf; + if (lane < 16) + { + float4 const mq + = *reinterpret_cast(red_m + ((h * QG + c) * 16 + lane) * 4); + mrow = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + } +#pragma unroll + for (int j = 0; j < 16; ++j) + { + float const m = __shfl_sync(0xffffffffu, mrow, j); + pa[k][j] = m == kNegInf ? 0.0F : ex2(v[k][j] - m); + } + if (tid == 0) + { + M3FA_CLK(41); + } + uint32_t pw[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) + { + __half2 const hp = __floats2half2_rn(pa[k][2 * j], pa[k][2 * j + 1]); + pw[j] = *reinterpret_cast(&hp); + } + // P^T as the MN-major B operand: 8 query rows of one token are one 16-byte core-matrix row. + *reinterpret_cast(pdst + (2 * c) * 2048) = make_uint4(pw[0], pw[1], pw[2], pw[3]); + *reinterpret_cast(pdst + (2 * c + 1) * 2048) + = make_uint4(pw[4], pw[5], pw[6], pw[7]); + M3FA_PROF_T(18, 0); + } + } + if (tid == 0) + { + M3FA_CLK(42); + } + M3FA_FENCE_PROXY_ASYNC(); // P -> tensor core reads + if (tid == 0) + { + M3FA_CLK(43); + } + // PV starts as soon as P and the V tiles are in: thread 0 waits for the softmax warps and the + // producer warps (V); the row sums and statistics below overlap the MMAs. + if (warp == 0) + { + asm volatile("bar.sync 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + else + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + if (kWarpMma ? warp == 0 : tid == 0) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm_u = tmem; + if (kWarpMma ? elect_one() : true) + { +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + uint32_t const a0 = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const b0 = smem_u32(smem + S::kP + h * kRows * kPage * 2); +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + umma_f16(tm_u + (BPC + h) * kRows, umma_desc(a0 + ks * 4096, 2048, 128), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 1, 1), ks > 0); + } + } + umma_commit(mma_bar); + } + if (kWarpMma) + { + __syncwarp(); + } + } + M3FA_PROF_T(19, 0); +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + float const lw = warp_reduce_scatter16(pa[k], lane); + if ((lane & 1) == 0) + { + red_l[((h * QG + c) * 16 + rs_row) * 4 + wq] = lw; + } + } + } + group_bar(); + M3FA_PROF_T(20, 0); + if (wq == 0 && lane < 16) + { +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + int const r = (h * QG + c) * 16 + lane; + float4 const mq = *reinterpret_cast(red_m + r * 4); + float4 const lq = *reinterpret_cast(red_l + r * 4); + float const m = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + float const l = (lq.x + lq.y) + (lq.z + lq.w); + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[h * kRows + c * 16 + lane] = scale * vgs; + row_lse[h * kRows + c * 16 + lane] = lse; + } + } + } + } + __syncthreads(); // row scales and LSEs visible (PV was issued above) + M3FA_PROF(9); + if (warp < kComputeWarps && warp / 4 < kPairs) + { + mbar_wait_cta(mma_bar, 1); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + } + } + M3FA_PROF(3); + } + else + { + asm volatile("cp.async.wait_all;\n" ::: "memory"); + if (OPJ && M3FA_C1_OPISSUE == 2 && tid == 32) + { + op_issue(p, smem, static_cast(blockIdx.x)); + } +#if M3FA_TRIGGER_AT == 2 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (kAttnTrigger || (FQKV && M3FA_MEGA_TRIGGER == 1)) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + for (int r = tid; r < BPC * kRows; r += kThreads) + { + row_lse[r] = kNegInf; + row_scale[r] = 0.0F; + } + __syncthreads(); + } + + // ---------------- Push every partial row to the CTA that merges its head (CTA r merges heads r*BPC..) -------- + if constexpr (S::kRecvAlias) + { + // Peers' partials overlay this CTA's tiles: every CTA's PV MMAs have completed. + asm volatile("barrier.cluster.arrive.aligned;\n" ::: "memory"); + } + if (!(kEarlyClw && !S::kRecvAlias)) + { + asm volatile("barrier.cluster.wait.aligned;\n" ::: "memory"); // peers' mbarriers are initialized + } + M3FA_PROF(4); + if (OPJ && M3FA_C1_OPISSUE == 4 && tid == 32) + { + op_issue(p, smem, static_cast(blockIdx.x)); + } + uint32_t const recv_base = smem_u32(recv); + if (kRs2 && warp < 8) + { + // Split push: warp group rg pushes rows 8 rg .. 8 rg + 7 (row j goes to CTA j, which merges head j). + int const rg = warp / 4; + int const d = (warp % 4) * 32 + lane; + float v[8]; + if (any_valid) + { + tmem_ld8(tmem + ((warp % 4) * 32 << 16) + BPC * kRows + 8 * rg, v); + } + else + { +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] = 0.0F; + } + } +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] *= row_scale[8 * rg + j2]; + } +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + uint32_t const dst_cta = 8 * rg + j2; + uint32_t const row_addr = recv_base + static_cast((cr * kOPitch + d) * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr, dst_cta)), + "r"(__float_as_uint(v[j2])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + else if (!kRs2 && warp < kComputeWarps && warp / 4 < kPairs) + { + int const g = warp / 4; + int const d = (warp % 4) * 32 + lane; // output dim +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + int const chunk = cr * BPC + h; + float v[16]; + if (any_valid && kRegCore) + { + float const* obuf + = reinterpret_cast(smem + S::kK16 + h * S::kTileBytes + S::kObufOff); +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] = obuf[(c * 16 + j) * kOPitch + d]; + } + } + else if (any_valid) + { + tmem_ld16(tmem + ((warp % 4) * 32 << 16) + (BPC + h) * kRows + c * 16, v); + } + else + { +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] = 0.0F; + } + } +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] *= row_scale[h * kRows + c * 16 + j]; + } + // One scalar store per row: a vector store's lanes would address several destination CTAs at once, + // which measured several times slower. +#pragma unroll + for (int j = 0; j < 16; ++j) + { + uint32_t const dst_cta = j / BPC; + uint32_t const row_addr + = recv_base + static_cast((((chunk * QG + c) * BPC + j % BPC) * kOPitch + d) * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr, dst_cta)), + "r"(__float_as_uint(v[j])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + } + } + else if (tid >= kComputeThreads) + { + // Every partial row's LSE (BPC * kRows of them; more than the producer threads for the widest geometries). + for (int r = tid - kComputeThreads; r < BPC * kRows; r += kThreads - kComputeThreads) + { + int const h = r / kRows; + int const c = (r % kRows) / 16; + int const j = r % 16; + uint32_t const dst_cta = j / BPC; + uint32_t const row_addr + = recv_base + static_cast((((cr * BPC + h) * QG + c) * BPC + j % BPC) * kOPitch * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr + kHeadDim * 4, dst_cta)), + "r"(__float_as_uint(row_lse[r])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + M3FA_PROF(10); + if (OPJ && M3FA_C1_OPISSUE == 5 && tid == 32) + { + op_issue(p, smem, static_cast(blockIdx.x)); + } +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + if (wr_kind[k] == 1) + { + *reinterpret_cast(wr_ptr[k]) = static_cast(wr_val[k] & 0xFFFFu); + if (lane % 4 == 0) + { + *wr_sptr[k] = static_cast(wr_val[k] >> 16); + } + } + else if (wr_kind[k] == 2) + { + *reinterpret_cast(wr_ptr[k]) = wr_val[k]; + } + } + { + uint32_t ready = 0; +#ifdef M3FA_DBG_HANG + long long spins = 0; +#endif + while (!ready) + { +#ifdef M3FA_DBG_HANG + if (++spins == (1LL << 24)) + { + printf("M3FA hang: block %d thread %d merge recv\n", blockIdx.x, threadIdx.x); + } +#endif + 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}\n" + : "=r"(ready) + : "r"(mbar) + : "memory"); + } + } + M3FA_PROF(5); + + // ---------------- Merge heads [cr * BPC, cr * BPC + BPC) (split-K merge kernel's LSE weighting) ---------------- + constexpr int kMergeRows = QG * BPC; // (token, local head) pairs, 16 chunk weights each + // Wide geometries: chunk weights once per merged row (into the softmax scratch, free by now), then the weighted + // sums. Narrow ones recompute the weights per output element: no extra barrier on the critical path. + constexpr bool kRowWeights = kMergeRows >= 4; + float* wmerge = red_m; // [kMergeRows][kBlocks]; red_m holds BPC * kRows * 4 = kMergeRows * 64 floats + [[maybe_unused]] auto chunk_weights = [&](int mr, float(&lc)[kBlocks]) + { + float lmax = kNegInf; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + lc[c] = recv[(c * kMergeRows + mr) * kOPitch + kHeadDim]; + lmax = fmaxf(lmax, lc[c]); + } + lmax = lmax == kNegInf ? 0.0F : lmax; + float denom = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + lc[c] = ex2(lc[c] - lmax); // now the chunk weight; 0 for an empty chunk + denom += lc[c]; + } + denom = denom > 0.0F ? denom : 1.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + // -0.0 marks an empty chunk (skipped below, as the merge kernel skips zero weights). + lc[c] = lc[c] != 0.0F ? __fdividef(lc[c], denom) : -0.0F; + } + }; + if constexpr (kRowWeights) + { + for (int mr = tid; mr < kMergeRows; mr += kThreads) + { + float lc[kBlocks]; + chunk_weights(mr, lc); +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wmerge[mr * kBlocks + c] = lc[c]; + } + } + __syncthreads(); + } +#pragma unroll + for (int idx = tid; idx < kMergeRows * kHeadDim; idx += kThreads) + { + int const mr = idx >> 7; + int const d = idx & (kHeadDim - 1); + float wn[kBlocks]; + if constexpr (kRowWeights) + { +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wn[c] = wmerge[mr * kBlocks + c]; + } + } + else + { +#ifndef M3FA_MEGA_MERGE_LOCAL +#define M3FA_MEGA_MERGE_LOCAL \ + 0 // 1 (measured +0.6 us slower): every merging thread derives the 16 chunk weights itself (broadcast shared loads, + // independent exps): no shuffle chain; the same arithmetic as chunk_weights +#endif + if constexpr (M3FA_MEGA_MERGE_LOCAL != 0) + { + chunk_weights(mr, wn); + } + else + { + // Narrow geometries: the warp (32 dims of one merged row) derives the 16 chunk weights together, lane c + // holding chunk c's, and broadcasts them. + float lc = (lane < kBlocks) ? recv[(lane * kMergeRows + mr) * kOPitch + kHeadDim] : kNegInf; + float lmax = lc; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + lmax = fmaxf(lmax, __shfl_xor_sync(0xffffffffu, lmax, o)); + } + lmax = lmax == kNegInf ? 0.0F : lmax; + float const wc = lane < kBlocks ? ex2(lc - lmax) : 0.0F; + // The denominator in chunk order from 0.0, as chunk_weights (and the merge kernel) add it: a butterfly + // reduction rounds differently in the last bit. + float denom = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + denom += __shfl_sync(0xffffffffu, wc, c); + } + denom = denom > 0.0F ? denom : 1.0F; + float const wl = wc != 0.0F ? __fdividef(wc, denom) : -0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wn[c] = __shfl_sync(0xffffffffu, wl, c); + } + } + } + if (tid == 0) + { + M3FA_CLK(45); + } + float acc = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + float const prod = recv[(c * kMergeRows + mr) * kOPitch + d] * wn[c]; + acc = __float_as_uint(wn[c]) != 0x80000000u ? acc + prod : acc; + } + if (tid == 0) + { + M3FA_CLK(46); + } + int const tl = mr / BPC; + int const head = cr * BPC + mr % BPC; + __nv_bfloat16 const ob = __float2bfloat16_rn(acc); + int const out_row = tok0 + tfirst + tl; + p.out[static_cast(out_row) * p.out_stride + head * kHeadDim + d] = ob; + if constexpr (MXFP8) + { + float const x = __bfloat162float(ob); + float amax = fabsf(x); +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, mask)); + } + float const sfv = amax * rcp_approx_ftz(448.0F); + int const bits = __float_as_int(sfv); + int ex = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0 ? 1 : 0); + ex = min(ex, 254); + float scale = __int_as_float((254 - ex) << 23); + scale = amax != 0.0F ? scale : 0.0F; + __nv_fp8_e4m3 const q8v(x * scale); + int64_t const q8off = static_cast(out_row) * (p.num_sf_k_tiles * 128) + head * kHeadDim + d; + if constexpr (FQKV || OPJ) + { + // Megakernel hand-off to the o_proj units (Lamport): this call's parity buffer, the other one reset to + // 0xFF. 0xFF is never written: a NaN code is canonicalized to 0x7F; the UE8M0 code is at most 254. + uint8_t const code = q8v.__x == 0xFFu ? static_cast(0x7Fu) : q8v.__x; + int64_t const cur = (mega_calls & 1u) * kMaxTok * kOpK; + int64_t const nxt = ((mega_calls & 1u) ^ 1u) * kMaxTok * kOpK; + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.q8 + cur + q8off), "h"(static_cast(code)) + : "memory"); + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.q8 + nxt + q8off), "h"(static_cast(0xFFu)) + : "memory"); + } + else + { + p.q8[q8off] = q8v.__x; + } + if (lane == 0) + { + int const k_sf = head * (kHeadDim / 32) + (d >> 5); + int64_t const off = static_cast(out_row / 128) * p.num_sf_k_tiles * 512 + (k_sf / 4) * 512 + + (out_row % 32) * 16 + ((out_row % 128) / 32) * 4 + k_sf % 4; + if constexpr (FQKV || OPJ) + { + int64_t const cur = (mega_calls & 1u) * kOpSfBytes; + int64_t const nxt = ((mega_calls & 1u) ^ 1u) * kOpSfBytes; + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.sf + cur + off), "h"(static_cast(ex)) + : "memory"); + asm volatile("st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.sf + nxt + off), + "h"(static_cast(0xFFu)) + : "memory"); + } + else + { + p.sf[off] = static_cast(ex); + } + } + } + } + M3FA_PROF(60); + if constexpr (OPJ) + { + // ---------------- c1: o_proj unit blockIdx.x (v52 m3_mx_gemv bit for bit, mega_v7's phase) ---------------- + if (any_valid) + { + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + __syncthreads(); // the merge's reads done; the K tiles are free for the activation overlay + int const op_u = static_cast(blockIdx.x); + if (M3FA_C1_OPISSUE == 3 && tid == 32) + { + op_issue(p, smem, op_u); + } + uint32_t const tm = *tmem_slot; + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if (elect_one()) + { + op_stream(smem, tm); + } + __syncwarp(); + } + op_act(p, smem, op_u, tid, mega_calls); + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + M3FA_PROF_T(36, 0); + __syncthreads(); + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if (elect_one()) + { + op_mma(smem, tm); + M3FA_PROF_T(50, threadIdx.x); + } + __syncwarp(); + } + if (warp >= kComputeWarps) + { + op_epilogue(p, smem, tm, tid, op_u); + } + M3FA_PROF_T(37, kComputeThreads); + __syncthreads(); +#if M3FA_C1_TRIGGER == 3 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // after this CTA's y rows / pushes +#endif + if (warp == kWarps - 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" ::"r"(tm), "r"(S::kTmemCols) + : "memory"); + } + if (tid == 0) + { + p.flags[32 + blockIdx.x] = mega_calls + 1u; + } + M3FA_PROF(6); + return; + } + if (any_valid) + { + // Every TMEM read completed above (tcgen05.wait::ld); release the columns. + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + __syncthreads(); + if (FQKV && tid == 0) + { + p.flags[32 + blockIdx.x] = mega_calls + 1u; + } + M3FA_PROF(61); + if (warp == kWarps - 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" ::"r"(*tmem_slot), "r"(S::kTmemCols) + : "memory"); + } + M3FA_PROF(6); +#else + __trap(); +#endif // MINIMAX_M3_ATTN_DEVICE_ENABLED +} + +template +cudaError_t launch_impl(Params const& prm, int num_reqs, bool pdl, cudaStream_t stream, int* max_clusters) +{ + auto kernel = m3_fused_decode_attn_kernel; + constexpr int smem = std::conditional_t, SmemX>::kBytes; + constexpr int kCs = kBlocks / BPC; + static bool const configured = [&]() + { + bool ok = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem) == cudaSuccess; + ok = ok + && (kCs <= 8 + || cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) == cudaSuccess); + if (!ok) + { + cudaGetLastError(); + } + return ok; + }(); + if (!configured) + { + return cudaErrorInvalidConfiguration; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(FQKV ? kQkvGrid : kCs * (QLEN / QG) * num_reqs); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cudaLaunchAttribute attr[2]; + attr[0].id = cudaLaunchAttributeClusterDimension; + attr[0].val.clusterDim.x = kCs; + attr[0].val.clusterDim.y = 1; + attr[0].val.clusterDim.z = 1; + attr[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[1].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + if (max_clusters != nullptr) + { + cfg.numAttrs = 1; + cudaError_t const err = cudaOccupancyMaxActiveClusters(max_clusters, kernel, &cfg); + if (err != cudaSuccess) + { + cudaGetLastError(); + } + return err; + } + cfg.numAttrs = 2; + return cudaLaunchKernelEx(&cfg, kernel, prm); +} + +// The one geometry the target launches: one request verifying 6 query tokens, 6 clusters of 16 CTAs (96 o_proj +// units). +cudaError_t dispatch(Params const& prm, int qlen, int qg, int bpc, int num_reqs, bool pdl, cudaStream_t stream) +{ + if (qlen == 6 && qg == 1 && bpc == 1) + { + return launch_impl<6, 1, 1, true, false, true>(prm, num_reqs, pdl, stream, nullptr); + } + return cudaErrorInvalidValue; +} + +static_assert(kMaxTok == minimax_m3_attn_c1::kMaxTokens && kOpN == minimax_m3_attn_c1::kOpN + && kOpK == minimax_m3_attn_c1::kOpK && kOpSfBytes == minimax_m3_attn_c1::kOpSfBytes + && kOpN / 128 * kOpSplits == minimax_m3_attn_c1::kCtas, + "header dimensions"); + +} // namespace + +namespace +{ +// TMA descriptor of a [n, k] E4M3 weight (128 x 128 boxes, 128-byte swizzle), cached per weight address. +bool weight_map(CUtensorMap& map, void const* w, int k, int n) +{ + static std::mutex mu; + static std::unordered_map cache; + std::lock_guard lock(mu); + auto it = cache.find(w); + if (it != cache.end()) + { + map = it->second; + return true; + } + static PFN_cuTensorMapEncodeTiled_v12000 encode = []() + { + void* fn = nullptr; + cudaDriverEntryPointQueryResult q; + cudaGetDriverEntryPointByVersion("cuTensorMapEncodeTiled", &fn, 12000, cudaEnableDefault, &q); + return reinterpret_cast(fn); + }(); + if (encode == nullptr) + { + return false; + } + cuuint64_t dims[2] = {static_cast(k), static_cast(n)}; + cuuint64_t strides[1] = {static_cast(k)}; + cuuint32_t box[2] = {128, 128}; + cuuint32_t elem[2] = {1, 1}; + CUresult const r = encode(&map, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(w), dims, strides, box, elem, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_L2_256B, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + if (r != CUDA_SUCCESS) + { + return false; + } + cache.emplace(w, map); + return true; +} +} // namespace + +cudaError_t invokeMinimaxM3AttnC1(MinimaxM3AttnC1Params& params, void const* oWeight, int qlen, int qg, int bpc, + int numReqs, bool pdl, cudaStream_t stream) +{ + if (!weight_map(params.wmap_o, oWeight, kOpK, kOpN)) + { + return cudaErrorInvalidValue; + } + return dispatch(params, qlen, qg, bpc, numReqs, pdl, stream); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/attnC1.h b/cpp/tensorrt_llm/kernels/minimaxM3/attnC1.h new file mode 100644 index 000000000000..b5f57301ed8e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/attnC1.h @@ -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. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace minimax_m3_attn_c1 +{ +constexpr int kCtas = 96; // attention CTAs == o_proj units (48 row tiles x 2 K splits) +constexpr int kMaxTokens = 8; // rows of each call-parity buffer of the MXFP8 hand-off +constexpr int kOpN = 6144; // o_proj output width +constexpr int kOpK = 2048; // o_proj input width (16 heads x 128) +constexpr int kOpSfBytes = 16 * 512; // one parity of the MXFP8 hand-off's swizzled scales +} // namespace minimax_m3_attn_c1 + +//! Kernel parameters of the attention + o_proj (c1) kernel (the device code's Params). +struct MinimaxM3AttnC1Params +{ + __nv_bfloat16 const* packed; + int64_t packed_stride; + uint8_t* kv_data; + int64_t d_page, d_plane; + uint8_t* kv_scale; + int64_t s_page, s_plane; + uint8_t* idx_k; + int64_t i_page, i_tok; + int64_t num_pages; + int const* out_cache_loc; + int const* position_ids; + int const* block_table; + int64_t bt_stride; + int bt_cols; + int const* seq_lens; + float const* kv_quant_scale; + float const* k_gs; + float const* v_gs; + __nv_bfloat16 const* q_w; + __nv_bfloat16 const* k_w; + __nv_bfloat16 const* ik_w; + float const* rope; + float eps; + float sm_scale; + __nv_bfloat16* out; + int64_t out_stride; + uint8_t* q8; + uint8_t* sf; + int num_sf_k_tiles; + unsigned long long* prof; // optional per-CTA clock64 probes (-DM3FA_PROFILE builds only) + // Fused QKV projection (FQKV kernels only): out[M, 2560] = mxfp8(x)[M, 6144] @ W^T. + __nv_bfloat16 const* x; // [M, 6144] layer input (the all-reduce output) + int ldx; + int num_tokens; // M + uint8_t const* wsf; // UE8M0 weight scales, 128x4-swizzled + float* part; // [2 parities][splits][8 tokens][2560] partial sums, sentinel at rest + int* counters; // [0]: call parity, [1]: CTAs done with this call (zero at rest) + int attn_clusters; // leading clusters that run the attention + uint8_t const* xq; // mega: the layer input's MXFP8 values [M, 6144] (row stride ldxq) ... + uint8_t const* xsf_act; // ... and its 128x4-swizzled UE8M0 scale atoms (atom k at byte 512 k) + int ldxq; + __nv_bfloat16* qkv_out; // mega: the packed QKV rows [2 parities][8][2560] the attention reads + uint8_t const* wsf_o; // mega: the o_proj weight's scale atoms + __nv_bfloat16* y_out; // mega: the o_proj output [M, 6144] + int* ar_ws; // mega, ar_mode: the next m3_mc_allreduce's multicast exchange (int32 words)... + int const* ar_flags; // ... its call flags (parity of a (token, column slice) slot = flags[t * 8 + slice] & 1) + int ar_rank; + int ar_world; + int ar_maxt; + int ar_mode; // 0: none; 2: push the o_proj rows into the exchange (as m3_mx_gemv's ar mode 2) + unsigned int* flags; // mega: [0, 20) tile counters (4 per call), [32 + CTA] the attention CTAs' targets + long long pf_a0, pf_b0, pf_a1, pf_b1; // mega: L2 prefetch regions issued in the attention window (bytes 0: none) + int pf_hint; // 1: evict_last policy on the prefetch + uint8_t const* wpre; // mega: optional pre-tiled QKV weight [N/128][K/128][16 KB swizzled smem image], or null + alignas(64) CUtensorMap wmap; // the [2560, 6144] E4M3 weight, 128 x 128 boxes, 128-byte swizzle + alignas(64) CUtensorMap wmap_o; // mega: the [6144, 2048] E4M3 o_proj weight, same boxes +}; + +//! Launch the c1 kernel on `stream`. The caller fills `params` except `wmap_o`, the o_proj weight's TMA +//! descriptor (cached per address). Only (qlen, qg, bpc) = (6, 1, 1) with one request is compiled; anything else, or +//! a failed TMA descriptor encode, returns cudaErrorInvalidValue. +cudaError_t invokeMinimaxM3AttnC1(MinimaxM3AttnC1Params& params, void const* oWeight, int qlen, int qg, int bpc, + int numReqs, bool pdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/attnMega.cu b/cpp/tensorrt_llm/kernels/minimaxM3/attnMega.cu new file mode 100644 index 000000000000..a4c8dd923f6f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/attnMega.cu @@ -0,0 +1,3673 @@ +/* + * 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 attention megakernel: QKV GEMV + fused decode attention + o_proj GEMV in one launch. + +#include "tensorrt_llm/kernels/minimaxM3/attnMega.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +// The kernel uses tcgen05/TMEM (including block-scaled MMAs, which family targets lack), cluster DSMEM and E2M1 +// conversions: its body is compiled for the arch-specific sm_100a / sm_103a targets only (see the minimax_m3_attn_src +// library in kernels/CMakeLists.txt). Any other target compiles a kernel that traps, so a launch on a build without +// those targets fails loudly instead of returning untouched outputs. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030) \ + && (defined(__CUDA_ARCH_SPECIFIC__) || defined(__CUDA_ARCH_FEAT_SM100_ALL) || defined(__CUDA_ARCH_FEAT_SM103_ALL)) +#define MINIMAX_M3_ATTN_DEVICE_ENABLED 1 +#else +#define MINIMAX_M3_ATTN_DEVICE_ENABLED 0 +#endif + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr int kHeadDim = 128; +constexpr int kRotDim = 64; +constexpr int kPage = 128; +constexpr int kHq = 16; +constexpr int kColK = kHq; // packed head slots: [Q x16 | K | V | idxQ | idxK] +constexpr int kColV = kHq + 1; +constexpr int kColIdxK = kHq + 3; +constexpr int kBlocks = 16; // KV blocks per request == topk == cluster size +constexpr int kComputeWarps = 16; // QK / softmax / PV and the dequant (one 16-byte chunk per thread) +constexpr int kWarps = kComputeWarps + 4; // plus producer-only warps, so new K/V/index-K never queue behind a Q head +constexpr int kThreads = kWarps * 32; +constexpr int kComputeThreads = kComputeWarps * 32; +constexpr int kOPitch = 132; // floats per partial-O row +constexpr float kNegInf = -__builtin_huge_valf(); + +using Params = MinimaxM3AttnMegaParams; + +// Optional per-CTA clock64 probes for development (-DM3FA_PROFILE and a prof tensor): [CTA][64] timestamps. +// -DM3FA_PROFILE_GT records %globaltimer (ns, comparable across SMs and kernels) instead of clock64. +#ifdef M3FA_PROFILE_GT +#define M3FA_PROFILE +#define M3FA_NOW(v) asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(v)) +#else +#define M3FA_NOW(v) v = clock64() +#endif +#ifdef M3FA_PROFILE +#define M3FA_PROF_T(i, thread) \ + do \ + { \ + if (p.prof != nullptr && threadIdx.x == (thread)) \ + { \ + unsigned long long now_; \ + M3FA_NOW(now_); \ + p.prof[blockIdx.x * 64 + (i)] = now_; \ + } \ + } while (0) +#else +#define M3FA_PROF_T(i, thread) \ + do \ + { \ + } while (0) +#endif +#define M3FA_PROF(i) M3FA_PROF_T(i, 0) +// clock64 probes (development): slot i of this CTA, written by every calling thread. +#ifdef M3FA_PROFILE +#define M3FA_CLK(i) \ + do \ + { \ + if (p.prof != nullptr) \ + { \ + p.prof[blockIdx.x * 64 + (i)] = clock64(); \ + } \ + } while (0) +#else +#define M3FA_CLK(i) \ + do \ + { \ + } while (0) +#endif + +#ifndef M3FA_MEGA_TRIGGER +// Megakernel PDL release: 0 o_proj CTAs after their activation (attention CTAs at exit); 1 as 0 with the attention +// CTAs after their KV reads; 2 every CTA right after its griddepcontrol.wait; 3 o_proj CTAs after their epilogue. +#define M3FA_MEGA_TRIGGER 3 +#endif +#ifndef M3FA_TRIGGER_AT +#define M3FA_TRIGGER_AT 2 // where the PDL dependents are released: 0 entry, 1 after the wait, 2 after the KV reads +#endif + +// Generic-proxy shared-memory writes (fp16 tiles, Q, P) become visible to the tensor core (async proxy). +#define M3FA_FENCE_PROXY_ASYNC() asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory") + +__device__ __forceinline__ void cp_async16(void* dst, void const* src) +{ + uint32_t const d = static_cast(__cvta_generic_to_shared(dst)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(d), "l"(src) : "memory"); +} + +__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__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +// One packed E2M1 byte -> half2 (low nibble in the low half). +__device__ __forceinline__ uint32_t e2m1x2_to_h2(uint32_t byte) +{ + uint32_t out; + asm("{\n .reg .b8 src;\n mov.b32 {src, _, _, _}, %1;\n cvt.rn.f16x2.e2m1x2 %0, src;\n}\n" : "=r"(out) : "r"(byte)); + return out; +} + +__device__ __forceinline__ uint32_t hmul2_u(uint32_t a, uint32_t b) +{ + __half2 r = __hmul2(*reinterpret_cast<__half2*>(&a), *reinterpret_cast<__half2*>(&b)); + return *reinterpret_cast(&r); +} + +__device__ __forceinline__ uint32_t e4m3_to_h2_splat(uint32_t byte) +{ + __half_raw const h = __nv_cvt_fp8_to_halfraw(static_cast<__nv_fp8_storage_t>(byte), __NV_E4M3); + return static_cast(h.x) | (static_cast(h.x) << 16); +} + +__device__ __forceinline__ uint32_t e4m3x2_to_h2(uint32_t two_bytes) +{ + __half2_raw const h = __nv_cvt_fp8x2_to_halfraw2(static_cast<__nv_fp8x2_storage_t>(two_bytes), __NV_E4M3); + return static_cast(h.x) | (static_cast(h.y) << 16); +} + +// ---- Exact copies of the C++ producer's helpers (fusedQKNormRopeKernel.cu / quantization.cuh) ---- +template +__device__ __forceinline__ uint32_t cvt_sf_group_shfl_mask() +{ + constexpr uint32_t groupSize = static_cast(NUM_THREADS_PER_SF); + constexpr uint32_t groupMask = 0xFFFFFFFFU >> (32U - groupSize); + uint32_t laneId = 0; + asm("mov.u32 %0, %%laneid;" : "=r"(laneId)); + return groupMask << (laneId & ~(groupSize - 1U)); +} + +__device__ __forceinline__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) +{ + 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"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]), "f"(array[4]), "f"(array[5]), "f"(array[6]), + "f"(array[7])); + return val; +} + +__device__ __forceinline__ uint16_t quantizeMinimaxM3Fp4x4( + float const (&elements)[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(elements[i]); + localMax = fmaxf(localMax, fabsf(__bfloat162float(rounded[i]))); + } + uint32_t const groupMask = cvt_sf_group_shfl_mask<4>(); + 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 quantValues[8] = {}; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + quantValues[i] = __bfloat162float(rounded[i]) * outputScale; + } + return static_cast(fp32_vec_to_e2m1(quantValues)); +} + +// RMSNorm (gemma 1+w) + NeoX RoPE over the first 64 dims, exactly as the producer: lane owns 4 dims. +__device__ __forceinline__ void norm_rope( + float (&e)[4], float const (&w)[4], float eps, float const (&cs)[4], float const (&sn)[4], int lane) +{ + float ss = 0.0F; + ss += e[0] * e[0] + e[1] * e[1]; + ss += e[2] * e[2] + e[3] * e[3]; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + ss = ss + __shfl_xor_sync(0xffffffffu, ss, mask, 32); + } + float const rmsReciprocal = rsqrtf(ss / static_cast(kHeadDim) + eps); +#pragma unroll + for (int i = 0; i < 4; ++i) + { + e[i] *= rmsReciprocal * (1.0F + w[i]); + } + __syncwarp(); + constexpr int kPairOffset = (kRotDim / 2) / 4; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + int const dim = lane * 4 + i; + float paired = __shfl_xor_sync(0xffffffffu, e[i], kPairOffset); + if (dim < kRotDim) + { + bool const firstHalf = dim < kRotDim / 2; + if (firstHalf) + { + paired = -paired; + } + e[i] = e[i] * cs[i] + paired * sn[i]; + } + } + __syncwarp(); +} + +__device__ __forceinline__ void load_w4(float (&w)[4], __nv_bfloat16 const* src, int lane) +{ + uint2 const raw = *reinterpret_cast(src + lane * 4); + float2 const a = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw.x)); + float2 const b = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw.y)); + w[0] = a.x; + w[1] = a.y; + w[2] = b.x; + w[3] = b.y; +} + +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ uint32_t mapa_shared(uint32_t addr, uint32_t rank) +{ + uint32_t r; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n" : "=r"(r) : "r"(addr), "r"(rank)); + return r; +} + +// QLEN query tokens per request; QG of them per cluster (QLEN / QG clusters per request); BPC KV blocks per CTA +// (cluster of kBlocks / BPC CTAs). +// Byte offset of element (r, k) in a UMMA no-swizzle K-major operand tile of 128 k-elements per row: 8x8 core +// matrices of 16-byte rows, 128 bytes apart along k (LBO) and 2048 bytes apart per 8 rows (SBO). The same formula +// with r = token, k = dim stores V so that it is the MN-major (dims contiguous) A operand of O^T = V^T P^T. +__device__ __forceinline__ int kmaj_off(int r, int k) +{ + return (r >> 3) * 2048 + (k >> 3) * 128 + (r & 7) * 16 + (k & 7) * 2; +} + +__device__ __forceinline__ uint64_t umma_desc(uint32_t addr, uint32_t lbo, uint32_t sbo) +{ + return static_cast((addr >> 4) & 0x3FFFu) | (static_cast((lbo >> 4) & 0x3FFFu) << 16) + | (static_cast((sbo >> 4) & 0x3FFFu) << 32) | (1ull << 46); +} + +// kind::f16 instruction descriptor: F32 accumulate, F16 operands, B K-major, M = 128. +__host__ __device__ constexpr uint32_t umma_idesc(int n, int a_mn_major, int b_mn_major) +{ + return (1u << 4) | (static_cast(a_mn_major) << 15) | (static_cast(b_mn_major) << 16) + | (static_cast(n >> 3) << 17) | (static_cast(128 >> 4) << 24); +} + +// 8-row version for the split softmax: the same lane-bit order (16, 8, 4, then 2 and 1 in full) as +// warp_reduce_scatter16, so every row is combined in the same association (fp add is commutative, not associative). +// After it lanes with bits 4..2 = (b4, b3, b2) hold row 4*b4 + 2*b3 + b2. +template +__device__ __forceinline__ float warp_reduce_scatter8(float (&a)[8], int lane) +{ +#pragma unroll + for (int step = 0; step < 3; ++step) + { + int const mask = 16 >> step; + int const half = 4 >> step; + bool const upper = (lane & mask) != 0; +#pragma unroll + for (int i = 0; i < half; ++i) + { + float const send = upper ? a[i] : a[i + half]; + float const keep = upper ? a[i + half] : a[i]; + float const recv = __shfl_xor_sync(0xffffffffu, send, mask); + a[i] = kMax ? fmaxf(keep, recv) : keep + recv; + } + } + float x = a[0]; + float const o2 = __shfl_xor_sync(0xffffffffu, x, 2); + x = kMax ? fmaxf(x, o2) : x + o2; + float const o1 = __shfl_xor_sync(0xffffffffu, x, 1); + return kMax ? fmaxf(x, o1) : x + o1; +} + +// Reduce 16 per-lane values over the warp, one value per row: after the reduce-scatter lane l holds row +// 8*b4 + 4*b3 + 2*b2 + b1 (b_k = bit k of l) reduced over all 32 lanes. 16 shuffles instead of 80. +template +__device__ __forceinline__ float warp_reduce_scatter16(float (&a)[16], int lane) +{ +#pragma unroll + for (int step = 0; step < 4; ++step) + { + int const mask = 16 >> step; + int const half = 8 >> step; + bool const upper = (lane & mask) != 0; +#pragma unroll + for (int i = 0; i < half; ++i) + { + float const send = upper ? a[i] : a[i + half]; + float const keep = upper ? a[i + half] : a[i]; + float const recv = __shfl_xor_sync(0xffffffffu, send, mask); + a[i] = kMax ? fmaxf(keep, recv) : keep + recv; + } + } + float const other = __shfl_xor_sync(0xffffffffu, a[0], 1); + return kMax ? fmaxf(a[0], other) : a[0] + other; +} + +__device__ __forceinline__ void umma_f16( + uint32_t tmem_d, uint64_t adesc, uint64_t bdesc, uint32_t idesc, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %4, 0;\n tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(idesc), "r"(accumulate) + : "memory"); +} + +__device__ __forceinline__ void umma_commit(uint32_t bar) +{ + asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];\n" ::"r"(bar) : "memory"); +} + +__device__ __forceinline__ void mbar_wait_cta(uint32_t bar, uint32_t parity) +{ + uint32_t ready = 0; +#ifdef M3FA_DBG_HANG + long long spins = 0; +#endif + while (!ready) + { +#ifdef M3FA_DBG_HANG + if (++spins == (1LL << 24)) + { + printf("M3FA hang: block %d thread %d mbar 0x%x parity %u\n", blockIdx.x, threadIdx.x, bar, parity); + } +#endif + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(bar), "r"(parity) + : "memory"); + } +} + +// Spin without suspending (v52's acc_done wait): a suspended try_wait wakes late on a tcgen05.commit arrival. +#ifndef M3FA_EPI_SPIN +#define M3FA_EPI_SPIN 1 +#endif +__device__ __forceinline__ void mbar_wait_epi(uint32_t bar, uint32_t parity) +{ + if (!M3FA_EPI_SPIN) + { + mbar_wait_cta(bar, parity); + return; + } + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(bar), "r"(parity) + : "memory"); + } +} + +__device__ __forceinline__ void tmem_ld8(uint32_t addr, float (&v)[8]) +{ + uint32_t r[8]; + asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]) + : "r"(addr)); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < 8; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +__device__ __forceinline__ void tmem_ld16(uint32_t addr, float (&v)[16]) +{ + uint32_t r[16]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, " + "%15}, [%16];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]) + : "r"(addr)); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < 16; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +// Warp-level fp16 MMA (m16n8k16, fp32 accumulate) and its shared-memory fragment loads. +__device__ __forceinline__ void ldsm_x4(uint32_t addr, uint32_t (&r)[4]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void ldsm_x2(uint32_t addr, uint32_t (&r)[2]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" : "=r"(r[0]), "=r"(r[1]) : "r"(addr)); +} + +__device__ __forceinline__ void ldsm_x2_trans(uint32_t addr, uint32_t (&r)[2]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0, %1}, [%2];\n" + : "=r"(r[0]), "=r"(r[1]) + : "r"(addr)); +} + +__device__ __forceinline__ void mma_f16_16816(float (&d)[4], uint32_t const (&a)[4], uint32_t const (&b)[2]) +{ + 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};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +// Register attention core (warp MMAs, no TMEM) for up to two tokens per cluster; the tcgen05 core above that. +// Off by default: on GB300 it measured ~0.25 us slower than the tcgen05 core at M=1 (GEMV + attention chain). +#ifndef M3FA_REG_CORE +#define M3FA_REG_CORE 0 +#endif +template +struct RegCore +{ + static constexpr bool kOn = M3FA_REG_CORE != 0 && QG <= 2; +}; + +template +struct Smem +{ + static constexpr int kRows = 16 * QG; // query rows (token, head) of this cluster = UMMA N + static constexpr int kTok = BPC * kPage; // block tokens per CTA + static constexpr int kTileBytes = kPage * kHeadDim * 2; // one fp16 K or V block tile in UMMA layout + static constexpr int kK16 = 0; + static constexpr int kV16 = kK16 + BPC * kTileBytes; + static constexpr int kQ = kV16 + BPC * kTileBytes; // Q rows, UMMA K-major [row][dim] + // P^T per block, UMMA MN-major [token][row], over the block's K tile: P is written after the QK MMAs (which + // read every K tile) have completed. + static constexpr int kP = kK16; + // Packed V tiles as copied from the cache: [block][128 rows x 64 data bytes | 1 KB of 4x4-quad scales]. + static constexpr int kRawV = kQ + kRows * kHeadDim * 2; + static constexpr int kRawVBlock = kPage * 64 + kPage * 8; + static constexpr int kRowScale = kRawV + BPC * kRawVBlock; + static constexpr int kRowLse = kRowScale + BPC * kRows * 4; + // Per-warp softmax partials [block][row][4 warps]: max, then sum. + static constexpr int kRedM = (kRowLse + BPC * kRows * 4 + 15) / 16 * 16; + static constexpr int kRedL = kRedM + BPC * kRows * 16; + // Partials pushed by every CTA of the cluster: [chunk][token][local head][kOPitch] floats, lse at column 128. + // When they do not fit beside the tiles they overlay the K/V tiles (and P, Q): peers then push only once every + // CTA of the cluster has completed its PV MMAs (a cluster barrier right before the push). + // Register core: per-warp row maxima and sums [block][row][warp of the block], two arrays. + static constexpr int kRedX = kRedL + BPC * kRows * 16; + static constexpr int kRedXBytes = RegCore::kOn ? 2 * BPC * kRows * (16 / BPC) * 4 : 0; + static constexpr int kStatsEnd = kRedX + kRedXBytes; + // Register core: P (fp16, [row][token] core matrices) at the start of its block's K tile, the scaled-free O + // partial ([row][kOPitch] fp32) behind it. + static constexpr int kObufOff = 8192; + static_assert(!RegCore::kOn || (kRows * kPage * 2 <= kObufOff && kObufOff + kRows * kOPitch * 4 <= kTileBytes), + "register-core P and O staging fit in a K tile"); + static constexpr int kRecvBytes = kBlocks * QG * BPC * kOPitch * 4; + static constexpr bool kRecvAlias = kStatsEnd + kRecvBytes + 32 > 227 * 1024; + static constexpr int kRecv = kRecvAlias ? 0 : kStatsEnd; + static constexpr int kMbar = kRecvAlias ? kStatsEnd : kStatsEnd + kRecvBytes; // merge, MMA mbarriers, TMEM addr + static constexpr int kBytes = kMbar + 32; + static constexpr int kTmemCols = BPC * 2 * kRows <= 32 ? 32 + : BPC * 2 * kRows <= 64 ? 64 + : BPC * 2 * kRows <= 128 ? 128 + : BPC * 2 * kRows <= 256 ? 256 + : 512; + static_assert(BPC * 2 * kRows <= 512, "TMEM columns"); + static_assert(!kRecvAlias || kRecvBytes <= kRawV, "aliased partials must fit in the tiles, Q and P"); + static_assert(kRows * kPage * 2 <= kTileBytes, "P^T fits in the K tiles"); + static_assert(kBytes <= 227 * 1024, "shared memory budget"); + static_assert(QLEN % QG == 0 && (BPC == 1 || BPC == 2), "unsupported geometry"); +}; + +// The packed NVFP4 bytes of one 16-byte chunk (32 dims of one block row) and its two block scales. +struct RawChunk +{ + uint4 data; + uint32_t scales; // byte 0: dims [0, 16), byte 1: dims [16, 32) of the chunk +}; + +// Chunk q of a block tile: lanes 8j..8j+7 of a warp take 8 consecutive rows of one 32-dim chunk, so every 16-byte +// fp16 store phase writes 8 distinct core-matrix rows (conflict-free in the UMMA layout). +__device__ __forceinline__ int chunk_row(int q) +{ + return (q >> 5) * 8 + (q & 7); +} + +__device__ __forceinline__ int chunk_col(int q) +{ + return (q >> 3) & 3; +} + +__device__ __forceinline__ RawChunk load_chunk( + uint8_t const* data_page, uint8_t const* scale_page, int row, int cc, bool is_v) +{ + RawChunk c; + c.data = *reinterpret_cast(data_page + row * 64 + cc * 16); + if (is_v) + { + // V scales: 4x4 token-quad order [row / 4][scale column][row % 4]. + uint8_t const* sc = scale_page + (row >> 2) * 32 + (row & 3); + c.scales = static_cast(sc[(cc * 2) * 4]) | (static_cast(sc[(cc * 2 + 1) * 4]) << 8); + } + else + { + c.scales = *reinterpret_cast(scale_page + row * 8 + cc * 2); + } + return c; +} + +__device__ __forceinline__ void store_chunk_fp16(uint8_t* tile, int row, int cc, RawChunk const& c, bool valid) +{ + uint32_t const s0 = e4m3_to_h2_splat(c.scales & 0xFFu); + uint32_t const s1 = e4m3_to_h2_splat((c.scales >> 8) & 0xFFu); + uint32_t const wd[4] = {c.data.x, c.data.y, c.data.z, c.data.w}; + uint32_t out[16]; +#pragma unroll + for (int i = 0; i < 16; ++i) + { + uint32_t const b = (wd[i >> 2] >> (8 * (i & 3))) & 0xFFu; + out[i] = valid ? hmul2_u(e2m1x2_to_h2(b), i < 8 ? s0 : s1) : 0u; + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + *reinterpret_cast(tile + kmaj_off(row, cc * 32 + j * 8)) + = make_uint4(out[4 * j], out[4 * j + 1], out[4 * j + 2], out[4 * j + 3]); + } +} + +// ---------------- Fused QKV projection: block-scaled MXFP8 on tcgen05 (N=2560, K=6144) ---------------- +// out[M, 2560] = mxfp8(x)[M, 6144] @ W^T, the computation of the CUTLASS MXFP8 GEMM: W is the A operand (128 output +// rows per MMA), the quantized activation the B operand (8 token columns), FP32 accumulation in TMEM. A unit is one +// (128-row tile, K split) pair; the 20 x kQkvSplits units run on the last CTAs of the grid. Before the PDL wait a +// CTA streams its unit's weight tiles and scale atoms through a shared-memory ring into TMEM; after it, it quantizes +// its x slice (mxfp8_quantize math), runs the MMAs and stores its FP32 partial rows. +// +// Hand-off without a flag: the partial buffers are double buffered by call parity and hold a sentinel (an FP32 NaN +// the MMA never produces) at rest, so a consumer polls the partials themselves. The producers of call n reset the +// other parity's buffers after their PDL wait; call n - 1, the last reader of those, is complete by then. +constexpr int kQkvN = 2560; +constexpr int kQkvK = 6144; +constexpr int kQkvKTiles = kQkvK / 128; +constexpr int kQkvRowTiles = kQkvN / 128; +#ifndef M3FA_QKV_SPLITS +#define M3FA_QKV_SPLITS 4 +#endif +#ifndef M3FA_QKV_RING +#define M3FA_QKV_RING kQkvMaxTiles // every tile of the unit in flight at once +#endif +#ifndef M3FA_QKV_NACC +#define M3FA_QKV_NACC 1 +#endif +#ifndef M3FA_QKV_SCALED_MMA +#define M3FA_QKV_SCALED_MMA 1 // 0: unscaled per-block MMAs with the scales applied in the epilogue (see below) +#endif +#ifndef M3FA_QKV_PIPE +// 1: the weight TMA is issued first at entry; the activation has its own region, copied right after the PDL wait; +// each weight tile's MMAs issue as soon as it lands (v52's order). 0: tiles into TMEM first, activation over the ring. +#define M3FA_QKV_PIPE 1 +#endif +#ifndef M3FA_QKV_A_SMEM +#define M3FA_QKV_A_SMEM 1 // scaled MMAs only; 1: they read the weight tiles from shared memory (no copy into TMEM) +#endif +// Default: block-scaled kind::mxf8f6f4 MMAs reading A (the weight tiles) in place from the ring, issued per tile as +// it lands (M3FA_QKV_PIPE), one accumulator; M3FA_QKV_A_SMEM=0 copies each tile into TMEM first. Alternative +// (M3FA_QKV_SCALED_MMA=0): one unscaled kind::f8f6f4 MMA per 32-wide K block into its own accumulator, A read in place +// from shared memory, the epilogue applying both UE8M0 scales. On GB300 an M128 x N8 x K32 MMA occupies the tensor pipe +// ~105 cycles block-scaled and ~60 unscaled, independent of N up to 64; the unscaled form's epilogue (320 accumulator +// columns) costs more than it saves. +constexpr bool kQkvScaled = M3FA_QKV_SCALED_MMA != 0; +constexpr bool kQkvASmem = !kQkvScaled || M3FA_QKV_A_SMEM != 0; +constexpr int kQkvSplits = M3FA_QKV_SPLITS; +constexpr int kQkvUnits = kQkvRowTiles * kQkvSplits; +constexpr int kQkvMaxTiles = (kQkvKTiles + kQkvSplits - 1) / kQkvSplits; +constexpr int kQkvRing = M3FA_QKV_RING; +constexpr int kQkvNacc = M3FA_QKV_NACC; +constexpr int kQkvGrid = 112; // 7 clusters of 16 (or 14 of 8): one wave with the attention +constexpr int kMaxTok = 8; +constexpr int kQkvWTile = 128 * 128; // E4M3 bytes of one 128 x 128 weight tile +constexpr int kQkvSfAtom = 512; // one 128-row x 4-block UE8M0 scale atom +constexpr uint32_t kQkvSentinel = 0xFFFFFFFFu; +constexpr int64_t kQkvPartStride = static_cast(kQkvSplits) * kMaxTok * kQkvN; // one parity +// TMEM columns of the one 512-column allocation of a CTA: attention [0, 64), then the GEMM's accumulators, scale +// atoms (A per tile, B per tile) and the A tiles (32 columns of 128 x 128 E4M3 each). +constexpr int kTmemAll = 512; +// Megakernel: the attention CTAs run no projection unit, so the GEMM starts at column 0 of a projection CTA. +constexpr int kTAcc = 0; +constexpr int kTSfa = kTAcc + 16; +constexpr int kTSfb = kTSfa + 4 * kQkvMaxTiles; +constexpr int kTA = (kTSfb + 4 * kQkvMaxTiles + 31) / 32 * 32; +static_assert(!kQkvScaled || kTA + 32 * kQkvMaxTiles <= kTmemAll, "TMEM columns"); +static_assert(kQkvScaled || kTAcc + 32 * kQkvMaxTiles <= kTmemAll, "TMEM columns (one accumulator per K block)"); +static_assert(!kQkvASmem || kQkvRing == kQkvMaxTiles, "shared-memory A needs every tile resident"); +static_assert(kQkvUnits <= kQkvGrid, "at most one unit per CTA"); +static_assert(kQkvNacc == 1 || kQkvNacc == 2 || kQkvNacc == 4, "accumulator count"); +static_assert(!M3FA_QKV_PIPE || kQkvScaled, "the pipelined projection uses the block-scaled MMAs"); +constexpr int kQkvOwnRows = 128 / kQkvSplits; // rows of a tile whose split-K sum a CTA owns +// Megakernel o_proj (v52 m3_mx_gemv for [6144, 2048]: split 2, contiguous k-tiles, one accumulator): unit +// o = CTA - 16 (the last 96 CTAs), row tile o / 2, split o % 2. +constexpr int kOpN = 6144; +constexpr int kOpK = 2048; +constexpr int kOpKTiles = kOpK / 128; +constexpr int kOpSplits = 2; +constexpr int kOpTiles = kOpKTiles / kOpSplits; +constexpr int kOpOwnRows = 128 / kOpSplits; +constexpr int kOpFirstCta = kQkvGrid - kOpN / 128 * kOpSplits; // 16 +constexpr int kOpSfBytes = kOpKTiles * 512; // one parity of the MXFP8 hand-off's scales +#ifndef M3FA_OP_L2PF +#define M3FA_OP_L2PF \ + 0 // 1: QKV CTAs prefetch their o_proj weight slice into L2 before the PDL wait (competes + // with the QKV weight stream for HBM behind the MoE) +#endif +// o_proj TMEM columns (after the QKV unit's MMAs completed): accumulator, A / B scale atoms, A (8 tiles x 32) +constexpr int kOAcc = 0; +constexpr int kOSfa = 16; +constexpr int kOSfb = kOSfa + 4 * kOpTiles; +constexpr int kOA = 128; +static_assert(kOA + 32 * kOpTiles <= kTmemAll && kOSfb + 4 * kOpTiles <= kOA, "o_proj TMEM columns"); +constexpr int kQkvMbStride = kMaxTok + 4; // mailbox row pitch (fp32 words) + +__host__ __device__ constexpr int qkv_tile0(int s) +{ + return s * kQkvKTiles / kQkvSplits; +} + +__device__ __forceinline__ int qkv_unit() +{ + int const u = kQkvGrid - 1 - static_cast(blockIdx.x); + return u < kQkvUnits ? u : -1; +} + +// Attention layout plus, for the fused projection, the weight ring in front (the attention tiles reuse that space +// once the ring has drained into TMEM); the activation tiles, the pre-wait V copy and the barriers sit behind. +template +struct SmemX : Smem +{ + using B = Smem; + static constexpr int kRingW = 0; // [ring][128 rows][128 B], 128-byte swizzle + static constexpr int kRingSf = kRingW + kQkvRing * kQkvWTile; // [ring][512 B] + static constexpr int kRingEnd = kRingSf + kQkvRing * kQkvSfAtom; + // The activation tiles overlay the ring: written after the PDL wait, once every weight tile went on into TMEM. + static constexpr int kXt = kRingW; // [tile][8][128 B] + static constexpr int kXsf = kXt + kQkvMaxTiles * 1024; // [tile][512 B] + static constexpr int kXsc = (kRingEnd + 127) / 128 * 128; + // Peers push into recv at any time (an empty CTA right away), so it stays clear of the ring. + static constexpr int kRecv = FQKV ? kXsc : B::kRecv; + static constexpr int kRawV = FQKV ? (kRecv + kBlocks * QG * BPC * kOPitch * 4 + 127) / 128 * 128 : B::kRawV; + // Pipelined projection: the activation tiles and scale atoms in their own region (a projection CTA runs no + // attention, so it overlays the attention's partial-O and V regions), 1024-byte aligned for the swizzle. + static constexpr int kXA = FQKV ? (kRecv + 1023) / 1024 * 1024 : 0; + static constexpr int kXASf = kXA + kQkvMaxTiles * 1024; + static constexpr int kXAEnd = kXASf + kQkvMaxTiles * kQkvSfAtom; + static constexpr int kMbarA = FQKV ? kRawV + BPC * B::kRawVBlock : B::kMbar; + static constexpr int kMbar = FQKV && M3FA_QKV_PIPE && kXAEnd > kMbarA ? (kXAEnd + 127) / 128 * 128 : kMbarA; + static constexpr int kGBar = kMbar + 32; // full[ring], empty[ring], a_ready, acc_done, tile_done[tiles] + static constexpr int kRedBar = kGBar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); // split-K partials in + static constexpr int kOBar = kRedBar + 8; // o_proj: full[8], a_ready, acc_done, reduce + static constexpr int kMailbox = (kOBar + 8 * (kOpTiles + 3) + 15) / 16 * 16; // [split][own row][pitch] fp32 + // o_proj weight ring = ring tiles [0, 8); its activation tiles, scale atoms and mailbox behind it. + static constexpr int kOXt = kRingW + kOpTiles * kQkvWTile; + static constexpr int kOXsf = kOXt + kOpTiles * 1024; + static constexpr int kOMailbox = kOXsf + kOpTiles * kQkvSfAtom; + static_assert(!FQKV || kOMailbox + kOpSplits * kOpOwnRows * kQkvMbStride * 4 <= kRingW + kQkvRing * kQkvWTile, + "o_proj activation and mailbox fit behind its ring"); + static constexpr int kBytes = FQKV ? kMailbox + kQkvSplits * kQkvOwnRows * kQkvMbStride * 4 : B::kBytes; + static constexpr int kTmemCols = FQKV ? kTmemAll : B::kTmemCols; + static constexpr bool kRecvAlias = !FQKV && B::kRecvAlias; + // The epilogue reads the weight scale atoms while the attention already writes its tiles. + static_assert(!FQKV || kQkvScaled || B::kRecv <= kRingSf, "attention regions overlap the scale atoms"); + // The other attention regions overlay the ring: a unit's CTA writes them only once the ring has drained. + static_assert(kBytes <= 227 * 1024, "shared memory budget"); +}; + +__device__ __forceinline__ void mbar_init(uint32_t bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(bar), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbar_expect_tx(uint32_t bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(bar), "r"(bytes) : "memory"); +} + +__device__ __forceinline__ void bulk_g2s(void* dst, void const* src, uint32_t bytes, uint32_t bar) +{ + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], " + "%4;\n" ::"r"(smem_u32(dst)), + "l"(src), "r"(bytes), "r"(bar), "l"(0x12F0000000000000ull) // EVICT_FIRST: weights are read once per layer + : "memory"); +} + +__device__ __forceinline__ void tma_load_2d(uint32_t dst, void const* map, int c0, int c1, uint32_t bar) +{ + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;\n" ::"r"(dst), + "l"(map), "r"(c0), "r"(c1), "r"(bar), "l"(0x12F0000000000000ull) + : "memory"); +} + +// K-major operand in the 128-byte swizzle: rows of 128 bytes, 8-row groups 1024 bytes apart. +__device__ __forceinline__ uint64_t umma_desc_sw128(uint32_t addr) +{ + return umma_desc(addr, 16, 1024) | (2ull << 61); +} + +__device__ __forceinline__ void tmem_cp_128x256b(uint32_t taddr, uint64_t desc) +{ + asm volatile("tcgen05.cp.cta_group::1.128x256b [%0], %1;\n" ::"r"(taddr), "l"(desc)); +} + +__device__ __forceinline__ void tmem_cp_sf(uint32_t taddr, uint32_t saddr) +{ + asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;\n" ::"r"(taddr), "l"(umma_desc(saddr, 16, 128))); +} + +// Block-scaled MXFP8 instruction descriptor: E4M3 x E4M3, UE8M0 scales, M = 128, N = 8, scale byte kb. +__host__ __device__ constexpr uint32_t qkv_idesc(int kb) +{ + return (static_cast(kb) << 4) | ((8u >> 3) << 17) | (1u << 23) | ((128u >> 7) << 27) + | (static_cast(kb) << 29); +} + +__device__ __forceinline__ void umma_mx( + uint32_t tmem_d, uint32_t tmem_a, uint64_t bdesc, uint32_t idesc, uint32_t sfa, uint32_t sfb, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" + " tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [%0], [%1], %2, %3, [%4], [%5], p;\n}\n" :: + "r"(tmem_d), + "r"(tmem_a), "l"(bdesc), "r"(idesc), "r"(sfa), "r"(sfb), "r"(accumulate)); +} + +__device__ __forceinline__ float pow2_from_ue8m0(uint32_t e) +{ + return __uint_as_float(e == 0 ? 0x00400000u : (e << 23)); +} + +// Unscaled E4M3 x E4M3 MMA, FP32 accumulator: M = 128, N = 8, K = 32, a fresh accumulator. +constexpr uint32_t kQkvIdescF8 = (1u << 4) | ((8u >> 3) << 17) | ((128u >> 4) << 24); + +__device__ __forceinline__ void umma_f8_ss(uint32_t tmem_d, uint64_t adesc, uint64_t bdesc) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, 0, 0;\n tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(kQkvIdescF8)); +} + +__device__ __forceinline__ void tmem_ld16_nowait(uint32_t addr, float* v) +{ + uint32_t r[16]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, " + "%15}, [%16];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]) + : "r"(addr)); +#pragma unroll + for (int i = 0; i < 16; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +__device__ __forceinline__ void umma_mx_ss( + uint32_t tmem_d, uint64_t adesc, uint64_t bdesc, uint32_t idesc, uint32_t sfa, uint32_t sfb, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" + " tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [%0], %1, %2, %3, [%4], [%5], p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(idesc), "r"(sfa), "r"(sfb), "r"(accumulate)); +} + +__device__ __forceinline__ bool elect_one() +{ + uint32_t pred; + asm volatile("{\n .reg .pred P;\n elect.sync _|P, 0xffffffff;\n selp.u32 %0, 1, 0, P;\n}\n" : "=r"(pred)); + return pred != 0; +} + +__device__ __forceinline__ int qkv_parity(Params const& p) +{ + int v; + asm volatile("ld.relaxed.gpu.global.s32 %0, [%1];\n" : "=r"(v) : "l"(p.counters) : "memory"); + return v & 1; +} + +// Barriers and TMEM (warp 1), then the first ring stages (warp 0). Warps 0 and 1 of every CTA call it. +template +__device__ __forceinline__ void qkv_setup( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int lane, int u) +{ + constexpr bool kPipe = M3FA_QKV_PIPE != 0; + if (kPipe && warp == 0 && lane == 0 && u >= 0) + { + // First thing at entry: this unit's whole weight slice. The issuing thread initializes the ring's full + // barriers itself; every waiter on them passes bar.sync 8 below, which this thread reaches after the init. + asm volatile("prefetch.tensormap [%0];\n" ::"l"(&p.wmap) : "memory"); + int const rt = u / kQkvSplits; + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + for (int i = 0; i < kQkvRing; ++i) + { + mbar_init(gbar + 8 * i, 1); + } + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + for (int i = 0; i < kQkvRing && i < nt; ++i) + { + uint32_t const full = gbar + 8 * i; + mbar_expect_tx(full, kQkvWTile + kQkvSfAtom); + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap, (kt0 + i) * 128, rt * 128, full); + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf + (static_cast(rt) * kQkvKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, full); + } + M3FA_PROF_T(34, 0); + } + if (warp == 1) + { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n" ::"r"(smem_u32(tmem_slot)), + "r"(kTmemAll) + : "memory"); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n" ::: "memory"); + M3FA_PROF_T(35, 32); + if (lane == 0) + { + for (int i = (kPipe && u >= 0) ? kQkvRing : 0; i < 2 * kQkvRing + 2 + kQkvMaxTiles; ++i) + { + // PIPE 2: tile_done[0] (unused by the scaled MMAs) is act_ready, one arrival per thread outside warp 1 + mbar_init(gbar + 8 * i, M3FA_QKV_PIPE == 2 && i == 2 * kQkvRing + 2 ? kThreads - 32 : 1); + } + // The split-K reduce barrier: this CTA's own arrival carries the peers' partial bytes. + uint32_t const red = gbar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); + mbar_init(red, 1); + mbar_expect_tx(red, (kQkvSplits - 1) * kQkvOwnRows * kMaxTok * 4); + uint32_t const ob = smem_u32(smem + SX::kOBar); + for (int i = 0; i < kOpTiles + 3; ++i) + { + mbar_init(ob + 8 * i, 1); + } + mbar_expect_tx(ob + 8 * (kOpTiles + 2), (kOpSplits - 1) * kOpOwnRows * kMaxTok * 4); + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + asm volatile("bar.sync 8, 64;\n" ::: "memory"); + if (!kPipe && warp == 0 && lane == 0 && u >= 0) + { + int const rt = u / kQkvSplits; + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + for (int i = 0; i < kQkvRing && i < nt; ++i) + { + uint32_t const full = gbar + 8 * i; + mbar_expect_tx(full, kQkvWTile + kQkvSfAtom); + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap, (kt0 + i) * 128, rt * 128, full); + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf + (static_cast(rt) * kQkvKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, full); + } + M3FA_PROF_T(34, 0); + } +} + +// The rest of the ring (warp 0, lane 0) and the copies into TMEM (warp 1, lane 0); before the PDL wait. +template +__device__ __forceinline__ void qkv_stream_lane( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int u) +{ + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + uint32_t const full = gbar; + uint32_t const a_ready = gbar + 16 * kQkvRing; + if (warp == 0 || !kQkvScaled) + { + return; // every tile was issued in qkv_setup; the unscaled MMAs read them in place + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + // Fully unrolled with compile-time offsets: tcgen05 operands stay in uniform registers (a per-instruction + // conversion of computed operands costs ~150 cycles per MMA or copy). + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint32_t const sf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + mbar_wait_cta(full + 8 * i, 0); + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if constexpr (!kQkvASmem) + { +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kTA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + } + tmem_cp_sf(tm + kTSfa + i * 4, sf0 + i * kQkvSfAtom); + } + } + M3FA_PROF_T(31, 32); + umma_commit(a_ready); +} + +__device__ __forceinline__ bool mbar_test(uint32_t bar, uint32_t parity) +{ + uint32_t ready; + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(bar), "r"(parity) + : "memory"); + return ready != 0; +} + +// PIPE 2, warp 1's elected thread from before the PDL wait to the last MMA. Tiles that land before the activation +// is ready go on into TMEM (A and scale atom), off the critical path; once the activation is in (act_ready), the +// MMAs issue in tile order: the copied tiles with A in TMEM, every later tile as it lands with A read in place from +// the ring. The MMA sequence (and so the accumulation order) is the same either way. +template +__device__ __forceinline__ void qkv_hybrid(Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int u) +{ + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + uint32_t const act_ready = gbar + 8 * (2 * kQkvRing + 2); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + SX::kXA)); + uint32_t const wsf0 = smem_u32(smem + SX::kRingSf); + uint32_t const xsf0 = smem_u32(smem + SX::kXASf); + int copied = 0; // tiles [0, copied) are in TMEM + bool act = false; +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt && !act) + { + while (!mbar_test(gbar + 8 * i, 0)) + { + if (mbar_test(act_ready, 0)) + { + act = true; + break; + } + } + if (!act) + { + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kTA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + copied = i + 1; + } + } + } + while (!act) + { + act = mbar_test(act_ready, 0); + } + M3FA_PROF_T(21, 32); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + bool const in_tmem = i < copied; + if (!in_tmem) + { + mbar_wait_epi(gbar + 8 * i, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + } + uint32_t const sfb = tm + kTSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + uint64_t const bd = b0 + ((i * 1024 + kb * 32) >> 4); + if (in_tmem) + { + umma_mx(tm + kTAcc, tm + kTA + j * 8, bd, qkv_idesc(kb), tm + kTSfa + i * 4, sfb, j > 0); + } + else + { + umma_mx_ss(tm + kTAcc, w0 + ((i * kQkvWTile + kb * 32) >> 4), bd, qkv_idesc(kb), tm + kTSfa + i * 4, + sfb, j > 0); + } + } + } + } + M3FA_PROF_T(22, 32); + M3FA_PROF_T(31, 32); + umma_commit(gbar + 16 * kQkvRing); // a_ready (for its waiters) + umma_commit(gbar + 16 * kQkvRing + 8); // acc_done +} + +// The other lanes of warps 0 and 1 hold at __syncwarp rather than at griddepcontrol.wait, which would stall the +// whole warp (lane 0's loop included) until the producer grid completes. +template +__device__ __forceinline__ void qkv_stream( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int lane, int u) +{ + if (M3FA_QKV_PIPE == 2 && u >= 0 && warp == 1) + { + if (elect_one()) + { + qkv_hybrid(p, smem, gbar, tmem_slot, u); + } + __syncwarp(); + return; + } + if (u < 0 || warp > 1 || M3FA_QKV_PIPE) + { + return; // pipelined: the MMA thread takes each tile as it lands + } + if (lane == 0) + { + qkv_stream_lane(p, smem, gbar, tmem_slot, warp, u); + } + __syncwarp(); +} + +// After the PDL wait. Every thread quantizes a piece of the x slice into the B tiles (named barrier 9); warp 1 then +// issues the scale copies and the MMAs. +template +__device__ __forceinline__ void qkv_mma( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int tid, int u) +{ + int const warp = tid >> 5; + int const s = u % kQkvSplits; + int const kt0 = qkv_tile0(s); + int const nt = qkv_tile0(s + 1) - kt0; + { + // The MXFP8 activation slice (E4M3 rows of the unit's k-tiles in 128-byte-swizzled B tiles, and their scale + // atoms), copied with plain 16-byte loads over the overlaid ring once its tiles went on into TMEM: every load + // first, then the stores (one L2 round trip). + constexpr int kXDst = M3FA_QKV_PIPE ? SX::kXA : SX::kXt; + constexpr int kXSfDst = M3FA_QKV_PIPE ? SX::kXASf : SX::kXsf; + if (!M3FA_QKV_PIPE) + { + mbar_wait_cta(gbar + 16 * kQkvRing, 0); // a_ready: the ring's weight tiles are in TMEM + } + constexpr int kXItems = kQkvMaxTiles * 8 * kMaxTok; // (tile, 16-byte chunk, token) + constexpr int kSfItems = kQkvMaxTiles * (kQkvSfAtom / 16); + constexpr bool kHyb = M3FA_QKV_PIPE == 2; + constexpr int kCopiers = kHyb ? kThreads - 32 : kThreads; // PIPE 2: warp 1 is issuing the tiles + if (kHyb && warp == 1) + { + return; + } + int const ctid = kHyb && tid >= 64 ? tid - 32 : tid; + constexpr int kPer = (kXItems + kSfItems + kCopiers - 1) / kCopiers; + uint4 v[kPer]; + int dst[kPer]; +#pragma unroll + for (int k = 0; k < kPer; ++k) + { + int const q = ctid + kCopiers * k; + dst[k] = -1; + if (q < kXItems) + { + int const t = q % kMaxTok; + int const c = (q / kMaxTok) % 8; + int const i = q / (kMaxTok * 8); + if (i < nt && t < p.num_tokens) + { + v[k] = *reinterpret_cast( + p.xq + static_cast(t) * p.ldxq + (kt0 + i) * 128 + c * 16); + dst[k] = kXDst + i * 1024 + t * 128 + ((c ^ t) * 16); + } + } + else if (q < kXItems + kSfItems) + { + int const j = q - kXItems; + int const i = j / (kQkvSfAtom / 16); + int const c = j % (kQkvSfAtom / 16); + if (i < nt) + { + v[k] = *reinterpret_cast( + p.xsf_act + static_cast(kt0 + i) * kQkvSfAtom + c * 16); + dst[k] = kXSfDst + i * kQkvSfAtom + c * 16; + } + } + } +#pragma unroll + for (int k = 0; k < kPer; ++k) + { + if (dst[k] >= 0) + { + *reinterpret_cast(smem + dst[k]) = v[k]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if (kHyb) + { + asm volatile("mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];\n" ::"r"(gbar + 8 * (2 * kQkvRing + 2)) + : "memory"); + return; + } + if (warp != 1) + { + asm volatile("bar.arrive 9, %0;\n" ::"n"(kThreads) : "memory"); + } + } + if (warp == 1) + { + M3FA_PROF_T(32, 32); + asm volatile("bar.sync 9, %0;\n" ::"n"(kThreads) : "memory"); + M3FA_PROF_T(21, 32); + // Operands computed by the whole warp (warp-uniform, kept in uniform registers); one elected lane issues. + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + (M3FA_QKV_PIPE ? SX::kXA : SX::kXt))); + uint64_t const a0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + if (elect_one()) + { + M3FA_CLK(41); + if constexpr (!kQkvScaled) + { +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + mbar_wait_cta(gbar + 8 * i, 0); // weight tile i landed + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + umma_f8_ss(tm + kTAcc + (i * 4 + kb) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4), + b0 + ((i * 1024 + kb * 32) >> 4)); + } + umma_commit(gbar + 8 * (2 * kQkvRing + 2 + i)); // tile i's blocks are in TMEM + } + } + } + else + { + uint32_t const xsf0 = smem_u32(smem + (M3FA_QKV_PIPE ? SX::kXASf : SX::kXsf)); + uint32_t const wsf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + if constexpr (M3FA_QKV_PIPE != 0) + { + // Tile i as it lands: its A copy and scale atom into TMEM, then its MMAs (issue order). + mbar_wait_cta(gbar + 8 * i, 0); + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + if (i == nt - 1) + { + M3FA_CLK(47); // the unit's last weight tile landed (clock64 builds) + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if constexpr (!kQkvASmem) + { +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b( + tm + kTA + (i * 4 + kb) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + } + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + } + uint32_t const sfb = tm + kTSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + uint64_t const bd = b0 + ((i * 1024 + kb * 32) >> 4); + if constexpr (kQkvASmem) + { + umma_mx_ss(tm + kTAcc + (j % kQkvNacc) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4), bd, + qkv_idesc(kb), tm + kTSfa + i * 4, sfb, j >= kQkvNacc); + } + else + { + umma_mx(tm + kTAcc + (j % kQkvNacc) * 8, tm + kTA + j * 8, bd, qkv_idesc(kb), + tm + kTSfa + i * 4, sfb, j >= kQkvNacc); + } + } + } + } + } + M3FA_CLK(42); + if (M3FA_QKV_PIPE) + { + M3FA_PROF_T(31, 32); + umma_commit(gbar + 16 * kQkvRing); // a_ready (kept for its waiters: every tile is consumed) + } + umma_commit(gbar + 16 * kQkvRing + 8); // acc_done + } + M3FA_PROF_T(22, 32); + } +} + +// Warps 16..19 (TMEM lanes 32 (warp % 4) + lane = the unit's rows): partial rows out, the other parity reset. +template +__device__ __forceinline__ void qkv_epilogue( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int tid, int u, int par, unsigned int calls) +{ + int const warp = tid >> 5; + int const lane = tid & 31; + if constexpr (kQkvScaled) + { + mbar_wait_epi(gbar + 16 * kQkvRing + 8, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + M3FA_PROF_T(24, 512); + if (tid == 512) + { + M3FA_CLK(43); + } + int const r = (warp & 3) * 32 + lane; // row within the unit's tile + uint32_t const taddr = *tmem_slot + (((warp & 3) * 32) << 16) + kTAcc; + float val[kMaxTok]; + if constexpr (kQkvScaled) + { + float v[32]; + tmem_ld16(taddr, *reinterpret_cast(v)); + if constexpr (kQkvNacc == 4) + { + tmem_ld16(taddr + 16, *reinterpret_cast(v + 16)); + } +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = kQkvNacc == 4 ? (v[t] + v[8 + t]) + (v[16 + t] + v[24 + t]) + : kQkvNacc == 2 ? v[t] + v[8 + t] + : v[t]; + } + } + else + { + // Sum over K blocks in order: block result x (weight scale x activation scale). + int const nt = qkv_tile0(u % kQkvSplits + 1) - qkv_tile0(u % kQkvSplits); + float const* xsc = reinterpret_cast(smem + SX::kXsc); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = 0.0f; + } +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + // Tile i's four blocks as soon as its MMAs complete: the epilogue trails the MMAs by one tile. + mbar_wait_cta(gbar + 8 * (2 * kQkvRing + 2 + i), 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + float d[32]; + tmem_ld16_nowait(taddr + i * 32, d); + tmem_ld16_nowait(taddr + i * 32 + 16, d + 16); + uint32_t const wsw = *reinterpret_cast( + smem + SX::kRingSf + i * kQkvSfAtom + (r % 32) * 16 + (r / 32) * 4); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + float const ws = pow2_from_ue8m0((wsw >> (8 * kb)) & 0xFFu); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = fmaf(d[kb * 8 + t], ws * xsc[(i * 4 + kb) * kMaxTok + t], val[t]); + } + } + } + } + } + // v52's split-K reduce: split s owns rows [32 s, 32 s + 32) of the tile; the other splits push their fp32 partials + // of those rows into its mailbox (st.async, completing its reduce barrier's bytes); the owner adds the splits in + // rank order from 0.0 and stores bf16. + int const s = u % kQkvSplits; + int const owner = r / kQkvOwnRows; + int const local = r % kQkvOwnRows; + uint32_t const red = gbar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); + uint32_t cta_rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + // Units are numbered down from the last CTA: split q of this tile is CTA rank cta_rank - (q - s). + if (owner != s) + { + uint32_t const peer = cta_rank + s - owner; + uint32_t const dst + = mapa_shared(smem_u32(smem + SX::kMailbox) + ((s * kQkvOwnRows + local) * kQkvMbStride) * 4, peer); + uint32_t const pbar = mapa_shared(red, peer); +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];\n" ::"r"( + dst + 16 * g), + "r"(__float_as_uint(val[4 * g])), "r"(__float_as_uint(val[4 * g + 1])), + "r"(__float_as_uint(val[4 * g + 2])), "r"(__float_as_uint(val[4 * g + 3])), "r"(pbar) + : "memory"); + } + } + else + { + M3FA_PROF_T(26, tid - lane); + // test_wait (no suspend): completions by peers' st.async bytes wake a suspended try_wait late. + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(red) + : "memory"); + } + M3FA_PROF_T(27, tid - lane); + float const* mb = reinterpret_cast(smem + SX::kMailbox); + int const row = (u / kQkvSplits) * 128 + r; + // Lamport hand-off: this call's rows go to buffer (calls & 1) with relaxed stores, the consumers poll the + // values themselves (0xFFFF, a NaN the rounding of a finite sum never produces, marks "not yet"); the other + // buffer, last read by the previous call (complete: this grid passed its PDL wait), is reset to the marker. + uint16_t* cur = reinterpret_cast(p.qkv_out) + (calls & 1u) * kMaxTok * kQkvN; + uint16_t* nxt = reinterpret_cast(p.qkv_out) + ((calls & 1u) ^ 1u) * kMaxTok * kQkvN; +#pragma unroll + float qsum[kMaxTok]; +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + float4 part[kQkvSplits]; +#pragma unroll + for (int q = 0; q < kQkvSplits; ++q) + { + part[q] = *reinterpret_cast(mb + (q * kQkvOwnRows + local) * kQkvMbStride + 4 * g); + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float sum = 0.0F; +#pragma unroll + for (int q = 0; q < kQkvSplits; ++q) + { + float const pj = j == 0 ? part[q].x : j == 1 ? part[q].y : j == 2 ? part[q].z : part[q].w; + sum = sum + (q == s ? val[4 * g + j] : pj); + } + qsum[4 * g + j] = sum; + } + } +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + float const sum = qsum[t]; + if (t < p.num_tokens) + { + __nv_bfloat16 const b = __float2bfloat16_rn(sum); + uint16_t bits = *reinterpret_cast(&b); + // Never the "not yet" marker: a 0xFFFF NaN sum is stored as the NaN 0x7FFF. Only a NaN projection row + // changes (its K/V cache bytes then differ from v52 + p6 in the NaN's encoding); finite rows are exact. + bits = bits == 0xFFFFu ? static_cast(0x7FFFu) : bits; +#ifndef M3FA_QKV_REL +#define M3FA_QKV_REL 0 // 1: the row store is a release (tried: whether it becomes visible to the pollers sooner) +#endif + if (M3FA_QKV_REL) + { + asm volatile( + "st.release.gpu.global.u16 [%0], %1;\n" ::"l"(cur + static_cast(t) * kQkvN + row), + "h"(bits) + : "memory"); + } + else + { + asm volatile( + "st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(cur + static_cast(t) * kQkvN + row), + "h"(bits) + : "memory"); + } + } + asm volatile("st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(nxt + static_cast(t) * kQkvN + row), + "h"(static_cast(0xFFFFu)) + : "memory"); + } + M3FA_PROF_T(28, tid - lane); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + M3FA_PROF_T(25, 512); + if (tid == 512) + { + M3FA_CLK(46); + } +} + +// ---------------- Megakernel o_proj (v52 m3_mx_gemv bit for bit: split 2, contiguous k-tiles, nacc 1) ---------------- +// The unit's weight tiles and scale atoms into the ring (warp 0, lane 0). From L2 when the CTA prefetched them there. +template +__device__ __forceinline__ void op_issue(Params const& p, uint8_t* smem, int ou) +{ + int const rt = ou / kOpSplits; + int const kt0 = (ou % kOpSplits) * kOpTiles; + uint32_t const ob = smem_u32(smem + SX::kOBar); + for (int i = 0; i < kOpTiles; ++i) + { + mbar_expect_tx(ob + 8 * i, kQkvWTile + kQkvSfAtom); + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap_o, (kt0 + i) * 128, rt * 128, ob + 8 * i); + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf_o + (static_cast(rt) * kOpKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, ob + 8 * i); + } +} + +// L2 prefetch of the unit's weight tiles and scale atoms (one lane; nothing waits for it). +// The QKV units' share of the window-1 L2 prefetch (the next consumer's weights), issued once their own QKV weight +// stream has landed (acc_done): the attention core reads little from HBM. +__device__ __forceinline__ void window_prefetch_l2(Params const& p, int u) +{ + constexpr long long kChunk = 32768; + long long const n0 = (p.pf_b0 + kChunk - 1) / kChunk, n1 = (p.pf_b1 + kChunk - 1) / kChunk; + uint64_t pol = 0; + if (p.pf_hint) + { + asm volatile("createpolicy.fractional.L2::evict_last.b64 %0, 1.0;\n" : "=l"(pol)); + } + for (long long c = u; c < n0 + n1; c += kQkvUnits) + { + long long const a = c < n0 ? p.pf_a0 + c * kChunk : p.pf_a1 + (c - n0) * kChunk; + long long const rem = c < n0 ? p.pf_b0 - c * kChunk : p.pf_b1 - (c - n0) * kChunk; + unsigned int const n = static_cast(rem < kChunk ? rem : kChunk) & ~15u; + if (n == 0) + { + continue; + } + if (p.pf_hint) + { + asm volatile("cp.async.bulk.prefetch.L2.global.L2::cache_hint [%0], %1, %2;\n" ::"l"(a), "r"(n), "l"(pol) + : "memory"); + } + else + { + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;\n" ::"l"(a), "r"(n) : "memory"); + } + } +} + +__device__ __forceinline__ void op_prefetch_l2(Params const& p, int ou) +{ + int const rt = ou / kOpSplits; + int const kt0 = (ou % kOpSplits) * kOpTiles; + for (int i = 0; i < kOpTiles; ++i) + { + asm volatile("cp.async.bulk.prefetch.tensor.2d.L2.global [%0, {%1, %2}];\n" ::"l"(&p.wmap_o), + "r"((kt0 + i) * 128), "r"(rt * 128) + : "memory"); + } + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;\n" ::"l"( + p.wsf_o + (static_cast(rt) * kOpKTiles + kt0) * kQkvSfAtom), + "r"(kOpTiles * kQkvSfAtom) + : "memory"); +} + +// The weight tiles into TMEM as they land (warp 1, elected lane, operands warp-uniform). +template +__device__ __forceinline__ void op_stream(uint8_t* smem, uint32_t tm) +{ + uint32_t const ob = smem_u32(smem + SX::kOBar); + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint32_t const sf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kOpTiles; ++i) + { + mbar_wait_cta(ob + 8 * i, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kOA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + tmem_cp_sf(tm + kOSfa + i * 4, sf0 + i * kQkvSfAtom); + } + umma_commit(ob + 8 * kOpTiles); // a_ready +} + +// The attention's MXFP8 rows of this unit's k-tiles (the heads [kt0, kt0 + 8)), Lamport-polled value by value: 0xFF +// is NaN in both E4M3 and UE8M0, which the merge never writes (see the merge). Every thread of the CTA, one item each. +template +__device__ __forceinline__ void op_act(Params const& p, uint8_t* smem, int ou, int tid, unsigned int calls) +{ + int const kt0 = (ou % kOpSplits) * kOpTiles; + uint8_t const* q8 = p.q8 + (calls & 1u) * kMaxTok * kOpK; + uint8_t const* sf = p.sf + (calls & 1u) * kOpSfBytes; + constexpr int kXItems = kOpTiles * 8 * kMaxTok; + constexpr int kSfItems = kOpTiles * (kQkvSfAtom / 16); + for (int q = tid; q < kXItems + kSfItems; q += kThreads) + { + if (q < kXItems) + { + int const t = q % kMaxTok; + int const c = (q / kMaxTok) % 8; + int const i = q / (kMaxTok * 8); + if (t < p.num_tokens) + { + uint4 v; + bool pending; + do + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(q8 + static_cast(t) * kOpK + (kt0 + i) * 128 + c * 16) + : "memory"); + uint32_t const w[4] = {v.x, v.y, v.z, v.w}; + pending = false; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + // a byte equal to 0xFF: (x & 0x7F7F7F7F) + 0x01010101 sets bit 7 of exactly the 0x7F/0xFF bytes + uint32_t const hit = (((w[j] & 0x7F7F7F7Fu) + 0x01010101u) & w[j] & 0x80808080u); + pending |= hit != 0u; + } + } while (pending); + *reinterpret_cast(smem + SX::kOXt + i * 1024 + t * 128 + ((c ^ t) * 16)) = v; + } + } + else if (q < kXItems + kSfItems) + { + int const j = q - kXItems; + int const i = j / (kQkvSfAtom / 16); + int const c = j % (kQkvSfAtom / 16); + uint8_t const* src = sf + (kt0 + i) * kQkvSfAtom + c * 16; + uint4 v; + if (c < p.num_tokens) + { + // Row c of the atom: its 4 block scales in bytes [16 c, 16 c + 4). + do + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(src) + : "memory"); + } while ((((v.x & 0x7F7F7F7Fu) + 0x01010101u) & v.x & 0x80808080u) != 0u); + } + else + { + v = *reinterpret_cast(src); + } + *reinterpret_cast(smem + SX::kOXsf + i * kQkvSfAtom + c * 16) = v; + } + } +} + +// The MMAs (warp 1, elected lane): v52's order (per k-tile: B scale copy, 4 MMAs; one accumulator). +template +__device__ __forceinline__ void op_mma(uint8_t* smem, uint32_t tm) +{ + uint32_t const ob = smem_u32(smem + SX::kOBar); + mbar_wait_cta(ob + 8 * kOpTiles, 0); // a_ready + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + SX::kOXt)); + uint32_t const xsf0 = smem_u32(smem + SX::kOXsf); +#pragma unroll + for (int i = 0; i < kOpTiles; ++i) + { + uint32_t const sfb = tm + kOSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + umma_mx(tm + kOAcc, tm + kOA + j * 8, b0 + ((i * 1024 + kb * 32) >> 4), qkv_idesc(kb), tm + kOSfa + i * 4, + sfb, j > 0); + } + } + umma_commit(ob + 8 * (kOpTiles + 1)); // acc_done +} + +// Warps 16..19: the pair's split-K reduce (owner rows [64 s, 64 s + 64)) and the bf16 rows of y. +template +__device__ __forceinline__ void op_epilogue(Params const& p, uint8_t* smem, uint32_t tm, int tid, int ou) +{ + int const warp = tid >> 5; + int const lane = tid & 31; + uint32_t const ob = smem_u32(smem + SX::kOBar); + // The store's kernel parameters, fetched before the wait: a first touch of their constant-bank line costs a miss. + __nv_bfloat16* y_out = p.y_out; + int num_tokens = p.num_tokens; + int ar_mode = p.ar_mode; + asm volatile("" : "+l"(y_out), "+r"(num_tokens), "+r"(ar_mode)); +#ifndef M3FA_OP_PARPRE +#define M3FA_OP_PARPRE 1 // read the all-reduce slot parities before the accumulator wait (a global load off the tail) +#endif + // The parity flags are stable here: the all-reduce that flips them last completed before this grid's PDL wait, + // and the next one reads them only after this kernel's pushes. + int par_pre[kMaxTok]; + { + int const row_pre = (ou / kOpSplits) * 128 + ((tid >> 5) & 3) * 32 + (tid & 31); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + par_pre[t] = 0; + if (M3FA_OP_PARPRE && ar_mode != 0 && t < num_tokens) + { + asm volatile("ld.volatile.global.s32 %0, [%1];\n" + : "=r"(par_pre[t]) + : "l"(p.ar_flags + t * 8 + row_pre / (kOpN / 8)) + : "memory"); + } + } + } + mbar_wait_epi(ob + 8 * (kOpTiles + 1), 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF_T(51, kComputeThreads); + int const r = (warp & 3) * 32 + lane; + float v[16]; + tmem_ld16(tm + (((warp & 3) * 32) << 16) + kOAcc, v); + int const s = ou % kOpSplits; + int const owner = r / kOpOwnRows; + int const local = r % kOpOwnRows; + uint32_t const red = ob + 8 * (kOpTiles + 2); + uint32_t cta_rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + if (owner != s) + { + uint32_t const peer = cta_rank + owner - s; // o_proj units count up with the CTA index + uint32_t const dst + = mapa_shared(smem_u32(smem + SX::kOMailbox) + ((s * kOpOwnRows + local) * kQkvMbStride) * 4, peer); + uint32_t const pbar = mapa_shared(red, peer); +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];\n" ::"r"( + dst + 16 * g), + "r"(__float_as_uint(v[4 * g])), "r"(__float_as_uint(v[4 * g + 1])), "r"(__float_as_uint(v[4 * g + 2])), + "r"(__float_as_uint(v[4 * g + 3])), "r"(pbar) + : "memory"); + } + } + else + { + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(red) + : "memory"); + } + M3FA_PROF_T(52, tid - lane); + float const* mb = reinterpret_cast(smem + SX::kOMailbox); + int const row = (ou / kOpSplits) * 128 + r; + float sums[kMaxTok]; +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + float4 part[kOpSplits]; +#pragma unroll + for (int q = 0; q < kOpSplits; ++q) + { + part[q] = *reinterpret_cast(mb + (q * kOpOwnRows + local) * kQkvMbStride + 4 * g); + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float sum = 0.0F; +#pragma unroll + for (int q = 0; q < kOpSplits; ++q) + { + float const pj = j == 0 ? part[q].x : j == 1 ? part[q].y : j == 2 ? part[q].z : part[q].w; + sum = sum + (q == s ? v[4 * g + j] : pj); + } + sums[4 * g + j] = sum; + } + } + M3FA_PROF_T(54, tid - lane); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + float const sum = sums[t]; +#ifndef M3FA_EXP +#define M3FA_EXP 0 +#endif + if (t < num_tokens) + { + if (M3FA_EXP == 1) + { + asm volatile( + "st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(y_out + static_cast(t) * kOpN + row), + "h"(__bfloat16_as_ushort(__float2bfloat16_rn(sum))) + : "memory"); + } + else + { + y_out[static_cast(t) * kOpN + row] = __float2bfloat16_rn(sum); + } + } + if (ar_mode != 0) + { + // m3_mx_gemv's push into the all-reduce exchange: lane pairs hold row pairs; the even lane writes the + // bf16x2 word into this rank's slot of the (token, column slice) parity the AR's flags select. + float const hi = __shfl_xor_sync(0xffffffffu, sum, 1); + if ((lane & 1) == 0 && t < num_tokens) + { + int par; + if (M3FA_OP_PARPRE) + { + par = par_pre[t]; + } + else + { + asm volatile("ld.volatile.global.s32 %0, [%1];\n" + : "=r"(par) + : "l"(p.ar_flags + t * 8 + row / (kOpN / 8)) + : "memory"); + } + int64_t const off = static_cast(par & 1) * p.ar_maxt * p.ar_world * (kOpN / 2) + + (static_cast(t) * p.ar_world + p.ar_rank) * (kOpN / 2) + row / 2; + __nv_bfloat162 w2 = __floats2bfloat162_rn(sum, hi); + uint32_t word = *reinterpret_cast(&w2); + // -0.0 halves to +0.0: 0x80000000 marks an empty exchange word (as m3_mx_gemv's _bf16x2_word) + word = (word & 0xFFFFu) == 0x8000u ? (word & 0xFFFF0000u) : word; + word = (word >> 16) == 0x8000u ? (word & 0x0000FFFFu) : word; + asm volatile("st.global.b32 [%0], %1;\n" ::"l"(p.ar_ws + off), "r"(word) : "memory"); + } + } + } + M3FA_PROF_T(53, tid - lane); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); +} + +// Packed-QKV element group [token][col * 128 + 4 * lane, +4): the split partials, polled until every one is stored, +// summed in split order and rounded to BF16. +__device__ __forceinline__ void qkv_load(Params const& p, int par, int tok, int col, int lane, float (&e)[4]) +{ + float const* base = p.part + par * kQkvPartStride + static_cast(tok) * kQkvN + col * kHeadDim + lane * 4; + uint4 pr[kQkvSplits]; + long long spins = 0; + bool ready; + do + { + ready = true; +#pragma unroll + for (int s = 0; s < kQkvSplits; ++s) + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(pr[s].x), "=r"(pr[s].y), "=r"(pr[s].z), "=r"(pr[s].w) + : "l"(base + static_cast(s) * kMaxTok * kQkvN) + : "memory"); + ready &= pr[s].x != kQkvSentinel && pr[s].y != kQkvSentinel && pr[s].z != kQkvSentinel + && pr[s].w != kQkvSentinel; + } + if (++spins > (1LL << 26)) + { + __trap(); + } + } while (!ready); + float s0 = __uint_as_float(pr[0].x), s1 = __uint_as_float(pr[0].y), s2 = __uint_as_float(pr[0].z), + s3 = __uint_as_float(pr[0].w); +#pragma unroll + for (int s = 1; s < kQkvSplits; ++s) + { + s0 += __uint_as_float(pr[s].x); + s1 += __uint_as_float(pr[s].y); + s2 += __uint_as_float(pr[s].z); + s3 += __uint_as_float(pr[s].w); + } + e[0] = __bfloat162float(__float2bfloat16_rn(s0)); + e[1] = __bfloat162float(__float2bfloat16_rn(s1)); + e[2] = __bfloat162float(__float2bfloat16_rn(s2)); + e[3] = __bfloat162float(__float2bfloat16_rn(s3)); +} + +// Every CTA once it has read the call parity and the partials it needs: the last one flips the parity. +__device__ __forceinline__ void qkv_done(Params const& p) +{ + if (atomicAdd(p.counters + 1, 1) == kQkvGrid - 1) + { + atomicExch(p.counters + 1, 0); + atomicXor(p.counters, 1); + } +} + +} // namespace + +namespace +{ + +// FQKV: the kernel also computes the QKV(+index) projection from the layer input (see qkv_mma) and runs on a fixed +// kQkvGrid-CTA grid; the first attn_clusters clusters then run the attention on the projection's output. +template +__global__ void __launch_bounds__(kThreads, 1) m3_fused_decode_attn_kernel(__grid_constant__ Params const p) +{ +#if MINIMAX_M3_ATTN_DEVICE_ENABLED + using S = SmemX; + constexpr int kRows = S::kRows; + constexpr int kTok = S::kTok; + constexpr int kCs = kBlocks / BPC; // cluster size + constexpr int kGroups = QLEN / QG; // clusters per request + constexpr uint32_t kPushBytes = (kHeadDim + 1) * 4; // one partial row: 128 outputs and the LSE + constexpr int kTasks = kHq * QG + 3 * QLEN; // producer tasks: 16 Q heads x QG, [K, V, index-K] x QLEN + constexpr int kMaxT = (kTasks + kWarps - 1) / kWarps; + constexpr int kChunks = BPC * kPage * 4; // 16-byte packed chunks per K (or V) tile set + constexpr int kKPer = (kChunks + kComputeThreads - 1) / kComputeThreads; // K chunks per thread [0, 512) + constexpr bool kRegCore = RegCore::kOn; // warp-MMA attention core (see RegCore) + // V tiles (raw copy before the PDL wait, dequantized during QK and the softmax): on the packed path every warp + // outside the softmax groups takes part (at M=1, warps 4..19), on the fused-projection path the producer warps. +#ifndef M3FA_RS2 +#define M3FA_RS2 1 // megakernel at one 16-row chunk: the softmax and push rows split over two warp groups (8 each) +#endif + constexpr bool kRs2 = FQKV && M3FA_RS2 && BPC * QG == 1 && !RegCore::kOn; + constexpr int kSoftmaxGroups = kRs2 ? 2 : (BPC * QG) < 4 ? (BPC * QG) : 4; +#ifndef M3FA_MEGA_VWIDE +#define M3FA_MEGA_VWIDE \ + 1 // megakernel attention CTAs (no projection unit): every warp outside the softmax copies and + // dequantizes V, as on the packed path (not only the four producer warps) +#endif + constexpr int kVStart = (kRegCore || (FQKV && !M3FA_MEGA_VWIDE)) ? kComputeThreads : 128 * kSoftmaxGroups; + constexpr int kVThreads = kThreads - kVStart; // V chunks: the warps outside the softmax groups + constexpr int kVPer = (kChunks + kVThreads - 1) / kVThreads; + constexpr int kPairs = BPC * QG; // (block, 16-row chunk) pairs of the softmax and push + constexpr int kPairsPerGroup = (kPairs + 3) / 4; // per group of four compute warps + // Named barrier 6: the softmax warps (P written) and the producer warps (V written) before the PV issue. + constexpr int kPvBarThreads = 128 * kSoftmaxGroups + kVThreads; + extern __shared__ __align__(128) uint8_t smem[]; + float* row_scale = reinterpret_cast(smem + S::kRowScale); + float* row_lse = reinterpret_cast(smem + S::kRowLse); + float* red_m = reinterpret_cast(smem + S::kRedM); + float* red_l = reinterpret_cast(smem + S::kRedL); + float* recv = reinterpret_cast(smem + S::kRecv); + uint32_t const mbar = smem_u32(smem + S::kMbar); + uint32_t const mma_bar = mbar + 8; + uint32_t* tmem_slot = reinterpret_cast(smem + S::kMbar + 16); + + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const cr = blockIdx.x % kCs; // cluster rank; owns blocks [cr * BPC, cr * BPC + BPC) + int const ci = blockIdx.x / kCs; + int const req = ci / kGroups; + int const tfirst = (ci % kGroups) * QG; // this cluster's query tokens [tfirst, tfirst + QG) of the request + int const tok0 = req * QLEN; + uint32_t const gbar = smem_u32(smem + S::kGBar); // FQKV: weight ring, A ready, accumulators done + int const qkv_u = FQKV ? qkv_unit() : -1; // FQKV: this CTA's projection unit, or -1 + + // ---------------- Prologue: nothing here depends on the producer kernel ---------------- +#if M3FA_TRIGGER_AT == 0 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (!FQKV) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + M3FA_PROF(0); +#ifdef M3FA_PROFILE + if (p.prof != nullptr && tid == 0) + { + uint32_t smid; + asm volatile("mov.u32 %0, %%smid;" : "=r"(smid)); + p.prof[blockIdx.x * 64 + 40] = smid; + } +#endif + if constexpr (FQKV) + { + // The projection weights stream into TMEM while the producer (the all-reduce) still runs. + if (warp < 2) + { + qkv_setup(p, smem, gbar, tmem_slot, warp, lane, qkv_u); + } + if (ci >= p.attn_clusters) + { + // A projection-only CTA. Every reduce barrier of the cluster is initialized before any peer pushes. + asm volatile("barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;\n" :: + : "memory"); + int const op_u = static_cast(blockIdx.x) - kOpFirstCta; // this CTA's o_proj unit (>= 0 here) +#ifndef M3FA_OP_EARLY +#define M3FA_OP_EARLY 0 // 1: o_proj-only CTAs stream their weights from launch (0: after the PDL wait) +#endif + if (M3FA_OP_EARLY && qkv_u < 0 && warp == 0 && lane == 0) + { + op_issue(p, smem, op_u); // an o_proj-only CTA: its weights stream from launch on + } + if (M3FA_OP_L2PF && qkv_u >= 0 && warp == 0 && lane == 0) + { + op_prefetch_l2(p, op_u); // into L2 during the all-reduce, behind this CTA's QKV ring + } + qkv_stream(p, smem, gbar, tmem_slot, warp, lane, qkv_u); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (M3FA_MEGA_TRIGGER == 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + // This CTA's call counter: read after the PDL wait (the previous call's CTA increments it at its end), and + // early, so its latency stays off the epilogue's critical path. + unsigned int const calls = p.flags[32 + blockIdx.x]; + M3FA_PROF(1); + if (!M3FA_OP_EARLY && qkv_u < 0 && warp == 0 && lane == 0) + { + op_issue(p, smem, op_u); + } + if (qkv_u >= 0) + { + int const par = 0; + qkv_mma(p, smem, gbar, tmem_slot, tid, qkv_u); + if (warp >= kComputeWarps) + { + qkv_epilogue(p, smem, gbar, tmem_slot, tid, qkv_u, par, calls); + } + if (warp == 0 && lane == 0) + { + // The ring (and the activation over it) is free once the QKV MMAs completed. + mbar_wait_cta(gbar + 16 * kQkvRing + 8, 0); + op_issue(p, smem, op_u); + if (p.pf_b0 > 0 || p.pf_b1 > 0) + { + window_prefetch_l2(p, qkv_u); + } + } + } + __syncthreads(); // the QKV accumulator reads are done: TMEM columns [0, 512) are the o_proj's + M3FA_PROF(2); + uint32_t const tm = *tmem_slot; + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if (elect_one()) + { + op_stream(smem, tm); + } + __syncwarp(); + } + op_act(p, smem, op_u, tid, calls); + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + M3FA_PROF_T(36, 0); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Late release: the dependents launch once every unit has its attention input (the rel_v4 rule). + if (M3FA_MEGA_TRIGGER < 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + __syncthreads(); + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if (elect_one()) + { + op_mma(smem, tm); + M3FA_PROF_T(50, threadIdx.x); + } + __syncwarp(); + } + if (warp >= kComputeWarps) + { + op_epilogue(p, smem, tm, tid, op_u); + } + M3FA_PROF_T(37, kComputeThreads); + __syncthreads(); + M3FA_PROF_T(55, 0); + if (M3FA_MEGA_TRIGGER == 3) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // after this CTA's y rows / pushes + } + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" ::"r"(*tmem_slot), "r"(kTmemAll) + : "memory"); + M3FA_PROF_T(56, 32); + } + if (tid == 0) + { + p.flags[32 + blockIdx.x] = calls + 1u; // a plain store: no load on the exit path + } + M3FA_PROF(39); +#ifndef M3FA_FINAL_CLBAR +#define M3FA_FINAL_CLBAR 0 +#endif + if (M3FA_FINAL_CLBAR) + { + // No peer may exit while its partials can still be pushed into it (DSMEM lifetime). + asm volatile("barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;\n" :: + : "memory"); + } + M3FA_PROF(38); + return; + } + } +#ifndef M3FA_MEGA_ATTN_WAIT_FIRST +#define M3FA_MEGA_ATTN_WAIT_FIRST 0 // 1: the attention CTAs read nothing (metadata, caches) before their PDL wait +#endif +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (FQKV && M3FA_MEGA_ATTN_WAIT_FIRST) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } +#endif + if (!FQKV && warp == kWarps - 1) // a producer warp: it has nothing else to do before the PDL wait + { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n" ::"r"(smem_u32(tmem_slot)), + "r"(S::kTmemCols) + : "memory"); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n" ::: "memory"); + } + int const seq_len = p.seq_lens[req]; + int page[BPC]; +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + int const b = cr * BPC + h; + page[h] = b < p.bt_cols ? p.block_table[static_cast(req) * p.bt_stride + b] : 0; + } + // New-token cache slots, for the tokens up to this cluster's last one. + int slot[QLEN]; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + slot[t] = t < tfirst + QG ? p.out_cache_loc[tok0 + t] : -1; + } + // This warp's producer tasks and, for every lane that applies RoPE, the task token's cos/sin. + float cs[kMaxT][4], sn[kMaxT][4]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + int const t = i < kHq * QG ? tfirst + i / kHq : (i - kHq * QG) / 3; + if (i < kTasks && lane < 16 && t < QLEN) + { + int64_t const row = static_cast(p.position_ids[tok0 + t]) * kRotDim; +#pragma unroll + for (int e = 0; e < 4; ++e) + { + int const coefficient = (lane * 4 + e) & (kRotDim / 2 - 1); + cs[k][e] = p.rope[row + coefficient]; + sn[k][e] = p.rope[row + kRotDim / 2 + coefficient]; + } + } + else + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + cs[k][e] = 1.0F; + sn[k][e] = 0.0F; + } + } + } + float qw[4], kw[4], ikw[4]; + load_w4(qw, p.q_w, lane); + load_w4(kw, p.k_w, lane); + load_w4(ikw, p.ik_w, lane); + float const kqs = p.kv_quant_scale[1]; + float const vqs = p.kv_quant_scale[2]; + float const sl2e = (p.sm_scale * 1.4426950409f) * p.k_gs[0]; + float const vgs = p.v_gs[0]; + // Longest attended length among this cluster's tokens, and each block's rows within it. + int const kv_max = max(seq_len - QLEN + tfirst + QG, 0); + int nvh[BPC]; +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + nvh[h] = min(max(kv_max - (cr * BPC + h) * kPage, 0), kPage); + } + bool const any_valid = nvh[0] > 0; + if (tid == 32) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(mbar), "r"(1)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(mma_bar), "r"(1)); + asm volatile( + "mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(mbar), "r"(kBlocks * QG * BPC * kPushBytes) + : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + } + // This thread's packed K / V chunks, held in registers until they are dequantized after the PDL wait. + RawChunk kraw[kKPer]; +#pragma unroll + for (int k = 0; k < kKPer; ++k) + { + int const q = tid + kComputeThreads * k; + kraw[k] = RawChunk{make_uint4(0u, 0u, 0u, 0u), 0u}; + int const h = q / (kPage * 4); + if (q < kChunks && tid < kComputeThreads && nvh[h] > 0) + { + kraw[k] = load_chunk(p.kv_data + static_cast(page[h]) * p.d_page, + p.kv_scale + static_cast(page[h]) * p.s_page, chunk_row(q % (kPage * 4)), + chunk_col(q % (kPage * 4)), false); + } + } + // V goes through cp.async into shared memory: an asynchronous copy does not hold the producer warps up the way + // an outstanding register load does, and they need V only after QK has been issued. + if (tid >= kVStart) + { +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + if (nvh[h] > 0) + { + uint8_t* dst = smem + S::kRawV + h * S::kRawVBlock; + uint8_t const* data = p.kv_data + static_cast(page[h]) * p.d_page + p.d_plane; + uint8_t const* scales = p.kv_scale + static_cast(page[h]) * p.s_page + p.s_plane; + for (int c = tid - kVStart; c < S::kRawVBlock / 16; c += kVThreads) + { + cp_async16(dst + c * 16, c < kPage * 4 ? data + c * 16 : scales + (c - kPage * 4) * 16); + } + } + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + if (warp == kWarps - 1) + { + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + // Every CTA's mbarrier is initialized before any peer pushes into it; the matching wait sits right before the + // push, so this barrier's latency hides behind the attention (with aliased partials it moves there). + if constexpr (!S::kRecvAlias) + { + asm volatile("barrier.cluster.arrive.relaxed.aligned;\n" ::: "memory"); + } + int kv_len[QG]; +#pragma unroll + for (int tl = 0; tl < QG; ++tl) + { + kv_len[tl] = max(seq_len - QLEN + tfirst + tl + 1, 0); + } + int const cta_tok0 = cr * kTok; // request position of this CTA's first block row + // New token t's row in this CTA's tiles (block * 128 + row), or -1: the CTA owns the block holding the token's + // position, and the token is attended by this cluster (t <= its last token). Every thread derives the same. + int new_row[QLEN]; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + int const pos = seq_len - QLEN + t; + int const blk = min(pos >> 7, kBlocks - 1); + int const spage = slot[t] / kPage; + bool const mine = t < tfirst + QG && pos >= 0 && blk / BPC == cr && slot[t] >= 0 && spage < p.num_pages; + new_row[t] = mine && spage == page[blk % BPC] ? (blk % BPC) * kPage + (slot[t] & (kPage - 1)) : -1; + } + auto is_new_row = [&](int r) + { + bool hit = false; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + hit |= new_row[t] == r; + } + return hit; + }; + // Cache writes of the new K/V/index-K are issued after this CTA has pushed its partials: only the fp16 tile rows + // are on the attention's path. kind 0: none, 1: NVFP4 data + scale, 2: FP8 index-K. + int wr_kind[kMaxT]; + uint8_t* wr_ptr[kMaxT]; + uint8_t* wr_sptr[kMaxT]; + uint32_t wr_val[kMaxT]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + wr_kind[k] = 0; + wr_ptr[k] = nullptr; + wr_sptr[k] = nullptr; + wr_val[k] = 0u; + } + uint32_t tmem = 0; + auto dequant_k = [&]() + { + if (tid < kComputeThreads) + { +#pragma unroll + for (int k = 0; k < kKPer; ++k) + { + int const q = tid + kComputeThreads * k; + if (q < kChunks) + { + int const h = q / (kPage * 4); + int const row = chunk_row(q % (kPage * 4)); + if (!is_new_row(h * kPage + row)) + { + store_chunk_fp16( + smem + S::kK16 + h * S::kTileBytes, row, chunk_col(q % (kPage * 4)), kraw[k], row < nvh[h]); + } + } + } + } + }; + // The producer warps' V tiles (raw tiles from cp.async, dequantized to fp16), then their PV-barrier arrival. + auto dequant_v = [&]() + { +#ifdef M3FA_DBG_NO_VDQ + if constexpr (!kRegCore) + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + return; +#endif + asm volatile("cp.async.wait_all;\n" ::: "memory"); + asm volatile("bar.sync 4, %0;\n" ::"n"(kVThreads) : "memory"); +#pragma unroll + for (int k = 0; k < kVPer; ++k) + { + int const q = tid - kVStart + kVThreads * k; + if (q < kChunks) + { + int const h = q / (kPage * 4); + int const row = chunk_row(q % (kPage * 4)); + int const cc = chunk_col(q % (kPage * 4)); + if (!is_new_row(h * kPage + row)) + { + uint8_t const* rawv = smem + S::kRawV + h * S::kRawVBlock; + uint8_t const* sc = rawv + kPage * 64 + (row >> 2) * 32 + (row & 3); + RawChunk c; + c.data = *reinterpret_cast(rawv + row * 64 + cc * 16); + c.scales + = static_cast(sc[(cc * 2) * 4]) | (static_cast(sc[(cc * 2 + 1) * 4]) << 8); + store_chunk_fp16(smem + S::kV16 + h * S::kTileBytes, row, cc, c, row < nvh[h]); + } + } + } + M3FA_FENCE_PROXY_ASYNC(); // V tiles -> tensor core reads (ordered by the barrier before the PV issue) + if constexpr (!kRegCore) + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + M3FA_PROF_T(16, kComputeThreads); + M3FA_PROF_T(17, kThreads - 1); + }; + if constexpr (FQKV) + { + qkv_stream(p, smem, gbar, tmem_slot, warp, lane, qkv_u); + } + M3FA_PROF(11); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (FQKV && M3FA_MEGA_TRIGGER == 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#if M3FA_TRIGGER_AT == 1 + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +#endif +#endif + // FQKV (megakernel): this call's QKV rows sit in buffer (calls & 1), polled value by value below. + unsigned int const mega_calls = FQKV ? p.flags[32 + blockIdx.x] : 0u; + __nv_bfloat16 const* qkv_src = FQKV ? p.packed + (mega_calls & 1u) * kMaxTok * kQkvN : p.packed; + M3FA_PROF(1); + int const par = 0; + if constexpr (FQKV) + { + if (qkv_u >= 0) + { + qkv_mma(p, smem, gbar, tmem_slot, tid, qkv_u); + } + } + + if constexpr (FQKV) + { + if (qkv_u >= 0) + { + if (warp >= kComputeWarps) + { + qkv_epilogue(p, smem, gbar, tmem_slot, tid, qkv_u, par, mega_calls); + } + else + { + // The attention regions overlay the weight ring: every weight byte has gone on into TMEM (or, with + // the MMAs reading the ring, every MMA has completed) first. + mbar_wait_cta(gbar + 16 * kQkvRing + (kQkvASmem ? 8 : 0), 0); + } + } + if (any_valid) + { + dequant_k(); + } +#ifndef M3FA_MEGA_EARLYV +#define M3FA_MEGA_EARLYV 1 // attention CTAs dequantize V here, while the QKV rows are still on their way +#endif +#ifndef M3FA_MEGA_EARLYCLW +#define M3FA_MEGA_EARLYCLW 1 // and pass the cluster barrier (peers' merge mbarriers initialized) here too +#endif + if (M3FA_MEGA_EARLYV && any_valid && tid >= kVStart) + { + dequant_v(); + } + if constexpr (M3FA_MEGA_EARLYCLW && !S::kRecvAlias) + { + asm volatile("barrier.cluster.wait.aligned;\n" ::: "memory"); + } + M3FA_PROF(14); + } + + if (any_valid) // CTA-uniform + { + // ---------------- Producer ---------------- + // Q rows of this cluster's tokens in every CTA. The new K/V of every token up to this cluster's last one + // whose position falls in this CTA's blocks go straight into the fp16 tiles; the CTA writes the caches, and + // the index-K, only for its own tokens. + uint2 raw[kMaxT]; + float fq[FQKV ? kMaxT : 1][4]; + bool act[kMaxT]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + raw[k] = make_uint2(0u, 0u); + act[k] = false; + if (i < kTasks) + { + int t, col; + bool active = true; + if (i >= kHq * QG) + { + int const kv = i - kHq * QG; + t = kv / 3; + int const kind = kv % 3; + col = kind == 0 ? kColK : (kind == 1 ? kColV : kColIdxK); + int const pos = seq_len - QLEN + t; + active = t < tfirst + QG && pos >= 0 && min(pos >> 7, kBlocks - 1) / BPC == cr && slot[t] >= 0 + && slot[t] / kPage < p.num_pages && (kind != 2 || t >= tfirst); + } + else + { + t = tfirst + i / kHq; + col = i % kHq; + } + act[k] = active; + if (active) + { + uint2 const* src = reinterpret_cast( + qkv_src + static_cast(tok0 + t) * p.packed_stride + col * kHeadDim + lane * 4); + if constexpr (FQKV) + { + uint2 v; + bool pending; + do + { + asm volatile("ld.relaxed.gpu.global.v2.b32 {%0, %1}, [%2];\n" + : "=r"(v.x), "=r"(v.y) + : "l"(src) + : "memory"); + pending = (v.x & 0xFFFFu) == 0xFFFFu || (v.x >> 16) == 0xFFFFu || (v.y & 0xFFFFu) == 0xFFFFu + || (v.y >> 16) == 0xFFFFu; + } while (pending); + raw[k] = v; + } + else + { + raw[k] = *src; + } + } + } + } + M3FA_PROF(15); + // While the packed-QKV loads are in flight: dequantize this thread's K chunks to fp16 (exact: e2m1 x e4m3 + // fits fp16). Rows past the longest attended length are zeroed, as the Triton producer masks its loads; the + // new tokens' rows are written by their producer tasks instead. (Building the tiles before the PDL wait was + // measured slower: the cache loads issued at entry are still in flight there, so the wait moves later.) + if constexpr (!FQKV) + { + dequant_k(); + } +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + if (i >= kTasks || !act[k]) + { + continue; + } + float e[4]; + { + float2 const a = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw[k].x)); + float2 const b = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw[k].y)); + e[0] = a.x; + e[1] = a.y; + e[2] = b.x; + e[3] = b.y; + } + if (i < kHq * QG) + { + norm_rope(e, qw, p.eps, cs[k], sn[k], lane); + __nv_fp8x2_e4m3 const lo(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const hi(make_float2(e[2], e[3])); + uint2 hv; + hv.x = e4m3x2_to_h2(lo.__x); + hv.y = e4m3x2_to_h2(hi.__x); + *reinterpret_cast(smem + S::kQ + kmaj_off(i, lane * 4)) = hv; // row i = token * 16 + head + continue; + } + int const t = (i - kHq * QG) / 3; + int const kind = (i - kHq * QG) % 3; + int const spage = slot[t] / kPage; + int const within = slot[t] & (kPage - 1); + if (kind == 2) + { + norm_rope(e, ikw, p.eps, cs[k], sn[k], lane); + __nv_fp8x2_e4m3 const lo(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const hi(make_float2(e[2], e[3])); + wr_kind[k] = 2; + wr_val[k] = static_cast(lo.__x) | (static_cast(hi.__x) << 16); + wr_ptr[k] = p.idx_k + static_cast(spage) * p.i_page + static_cast(within) * p.i_tok + + lane * 4; + continue; + } + bool const isV = kind == 1; + if (!isV) + { + norm_rope(e, kw, p.eps, cs[k], sn[k], lane); + } + uint8_t sbyte = 0; + uint16_t const fp4 = quantizeMinimaxM3Fp4x4(e, isV ? vqs : kqs, &sbyte); + if (t >= tfirst) + { + int const scale_col = lane / 4; + int const scale_within = isV ? (within / 4) * 32 + scale_col * 4 + within % 4 : within * 8 + scale_col; + wr_kind[k] = 1; + wr_val[k] = static_cast(fp4) | (static_cast(sbyte) << 16); + wr_ptr[k] = p.kv_data + static_cast(spage) * p.d_page + (isV ? p.d_plane : 0) + within * 64 + + lane * 2; + wr_sptr[k] = p.kv_scale + static_cast(spage) * p.s_page + (isV ? p.s_plane : 0) + scale_within; + } + if (new_row[t] >= 0) + { + // The attention reads this page: its fp16 tile row is the freshly quantized row, dequantized like + // the rest of the tile. + int const h = new_row[t] / kPage; + int const row = new_row[t] % kPage; + uint32_t const sc2 = e4m3_to_h2_splat(sbyte); + uint2 hv; + hv.x = hmul2_u(e2m1x2_to_h2(fp4 & 0xFFu), sc2); + hv.y = hmul2_u(e2m1x2_to_h2(fp4 >> 8), sc2); + *reinterpret_cast( + smem + (isV ? S::kV16 : S::kK16) + h * S::kTileBytes + kmaj_off(row, lane * 4)) + = hv; + } + } + // K tiles, Q rows and patched rows -> tensor core reads. Only threads that wrote fence: the fence also + // waits for the thread's outstanding loads, and the producer warps still have their V chunks in flight. + bool wrote = tid < kComputeThreads; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + wrote |= act[k]; + } + if (wrote) + { + M3FA_FENCE_PROXY_ASYNC(); + } + M3FA_PROF(13); + __syncthreads(); +#if M3FA_TRIGGER_AT == 2 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Every K/V byte this CTA reads has arrived (the V chunks sit in registers): release the dependents now so + // their weight streams do not queue ahead of the KV reads (the megakernel releases late, from its o_proj). + if constexpr (!FQKV || M3FA_MEGA_TRIGGER == 1) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + tmem = *tmem_slot; + M3FA_PROF(2); + // ---------------- S^T = K Q^T per block on the tensor core: TMEM lane = token, column = query row ---------- +#ifndef M3FA_MEGA_WARPMMA +#define M3FA_MEGA_WARPMMA 1 // attention MMAs issued by warp 0's elected lane with warp-uniform operands (not tid 0) +#endif + if (!kRegCore && (FQKV && M3FA_MEGA_WARPMMA ? warp == 0 : tid == 0)) + { + // Operands computed by the whole warp stay in uniform registers; one elected lane issues. + uint32_t const tm_u = tmem; + bool const issue = FQKV && M3FA_MEGA_WARPMMA ? elect_one() : true; + if (issue) + { +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + uint32_t const a0 = smem_u32(smem + S::kK16 + h * S::kTileBytes); + uint32_t const b0 = smem_u32(smem + S::kQ); +#pragma unroll + for (int ks = 0; ks < kHeadDim / 16; ++ks) + { + umma_f16(tm_u + h * kRows, umma_desc(a0 + ks * 256, 128, 2048), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 0, 0), ks > 0); + } + } + umma_commit(mma_bar); + } + if (FQKV && M3FA_MEGA_WARPMMA) + { + __syncwarp(); + } + } + if (tid >= kVStart && !(FQKV && M3FA_MEGA_EARLYV)) + { + // The producer warps dequantize V while QK and the softmax run. + dequant_v(); + } + if constexpr (kRegCore) + { + // ---------------- Register core: warp w of block h = w / kWpb owns 8 * kNt tokens for S and 8 * kNt + // output dims for O; m16n8k16 fp16 MMAs with fp32 accumulators, the softmax in the S fragments ---------- + constexpr int kWpb = kComputeWarps / BPC; // warps per block + constexpr int kNt = 16 / kWpb; // 8-token (S) or 8-dim (O) tiles per warp + float* redx_m = reinterpret_cast(smem + S::kRedX); + float* redx_l = redx_m + BPC * kRows * kWpb; + if (warp < kComputeWarps) + { + int const h = warp / kWpb; + int const wib = warp % kWpb; + int const g = lane >> 2; + int const q4 = lane & 3; + uint32_t const ktile = smem_u32(smem + S::kK16 + h * S::kTileBytes); + uint32_t const vtile = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const qrows = smem_u32(smem + S::kQ); + uint8_t* ptile = smem + S::kK16 + h * S::kTileBytes; // P over this block's K tile, after QK + auto block_bar = [&]() + { + if (h == 0) + { + asm volatile("bar.sync 14, %0;\n" ::"n"(kWpb * 32) : "memory"); + } + else + { + asm volatile("bar.sync 15, %0;\n" ::"n"(kWpb * 32) : "memory"); + } + }; + // S = Q K^T: rows (c, g / g + 8), tokens 8 tb + 2 q4 (+1). + float sv[QG][kNt][4]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + sv[c][nt][e] = 0.0F; + } + } + } +#pragma unroll + for (int ks = 0; ks < kHeadDim / 16; ++ks) + { + uint32_t bk[kNt][2]; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const tb = wib * kNt + nt; + ldsm_x2(ktile + tb * 2048 + (ks * 2 + ((lane >> 3) & 1)) * 128 + (lane & 7) * 16, bk[nt]); + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { + uint32_t aq[4]; + ldsm_x4( + qrows + (c * 2 + ((lane >> 3) & 1)) * 2048 + (ks * 2 + (lane >> 4)) * 128 + (lane & 7) * 16, + aq); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + mma_f16_16816(sv[c][nt], aq, bk[nt]); + } + } + } + M3FA_PROF(7); + // Row maxima (log2 domain, masked as the producer masks), quad then across the block's warps. + float mrow[QG][2]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { + float mx0 = kNegInf, mx1 = kNegInf; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + int const tok = cta_tok0 + h * kPage + (wib * kNt + nt) * 8 + 2 * q4 + (e & 1); + bool const valid = tok < kv_max && tok < kv_len[c]; + float const v = valid ? sv[c][nt][e] * sl2e : kNegInf; + sv[c][nt][e] = v; + if (e < 2) + { + mx0 = fmaxf(mx0, v); + } + else + { + mx1 = fmaxf(mx1, v); + } + } + } + mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, 1)); + mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, 2)); + mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, 1)); + mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, 2)); + if (q4 == 0) + { + redx_m[(h * kRows + c * 16 + g) * kWpb + wib] = mx0; + redx_m[(h * kRows + c * 16 + g + 8) * kWpb + wib] = mx1; + } + } + block_bar(); // every warp of the block is past QK (P may overwrite the K tile) and wrote its maxima + M3FA_PROF(8); +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int half = 0; half < 2; ++half) + { + float const* rm = redx_m + (h * kRows + c * 16 + g + 8 * half) * kWpb; + float m = kNegInf; +#pragma unroll + for (int w4 = 0; w4 < kWpb; w4 += 4) + { + float4 const v = *reinterpret_cast(rm + w4); + m = fmaxf(m, fmaxf(fmaxf(v.x, v.y), fmaxf(v.z, v.w))); + } + mrow[c][half] = m; + } + // P = exp2(s - m) as fp16 into [row][token] core matrices; the fp32 values feed the row sums. + float l0 = 0.0F, l1 = 0.0F; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + float a[4]; +#pragma unroll + for (int e = 0; e < 4; ++e) + { + float const m = mrow[c][e >> 1]; + a[e] = m == kNegInf ? 0.0F : ex2(sv[c][nt][e] - m); + } + l0 += a[0] + a[1]; + l1 += a[2] + a[3]; + int const t0 = (wib * kNt + nt) * 8 + 2 * q4; + int const r0 = c * 16 + g; + __half2 const p0 = __floats2half2_rn(a[0], a[1]); + __half2 const p1 = __floats2half2_rn(a[2], a[3]); + *reinterpret_cast<__half2*>( + ptile + (r0 >> 3) * 2048 + (t0 >> 3) * 128 + (r0 & 7) * 16 + (t0 & 7) * 2) + = p0; + *reinterpret_cast<__half2*>( + ptile + ((r0 + 8) >> 3) * 2048 + (t0 >> 3) * 128 + (r0 & 7) * 16 + (t0 & 7) * 2) + = p1; + } + l0 += __shfl_xor_sync(0xffffffffu, l0, 1); + l0 += __shfl_xor_sync(0xffffffffu, l0, 2); + l1 += __shfl_xor_sync(0xffffffffu, l1, 1); + l1 += __shfl_xor_sync(0xffffffffu, l1, 2); + if (q4 == 0) + { + redx_l[(h * kRows + c * 16 + g) * kWpb + wib] = l0; + redx_l[(h * kRows + c * 16 + g + 8) * kWpb + wib] = l1; + } + } + M3FA_PROF_T(19, 0); + } + __syncthreads(); // P, the V tiles (producer warps) and the row sums visible + M3FA_PROF(9); + if (warp < kComputeWarps) + { + int const h = warp / kWpb; + int const wib = warp % kWpb; + int const g = lane >> 2; + int const q4 = lane & 3; + uint32_t const vtile = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const ptile = smem_u32(smem + S::kK16 + h * S::kTileBytes); + float* obuf = reinterpret_cast(smem + S::kK16 + h * S::kTileBytes + S::kObufOff); + // O = P V: rows (c, g / g + 8), dims 8 dt + 2 q4 (+1). + float ov[QG][kNt][4]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + ov[c][nt][e] = 0.0F; + } + } + } +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + uint32_t bv[kNt][2]; + int const tok = ks * 16 + ((lane >> 3) & 1) * 8 + (lane & 7); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const dt = wib * kNt + nt; + ldsm_x2_trans(vtile + (tok >> 3) * 2048 + dt * 128 + (tok & 7) * 16, bv[nt]); + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { + uint32_t ap[4]; + ldsm_x4( + ptile + (c * 2 + ((lane >> 3) & 1)) * 2048 + (ks * 2 + (lane >> 4)) * 128 + (lane & 7) * 16, + ap); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + mma_f16_16816(ov[c][nt], ap, bv[nt]); + } + } + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const d = (wib * kNt + nt) * 8 + 2 * q4; + int const r0 = c * 16 + g; + *reinterpret_cast(obuf + r0 * kOPitch + d) = make_float2(ov[c][nt][0], ov[c][nt][1]); + *reinterpret_cast(obuf + (r0 + 8) * kOPitch + d) + = make_float2(ov[c][nt][2], ov[c][nt][3]); + } + } + } + else if (tid - kComputeThreads < BPC * kRows) + { + // Row statistics of every (block, row): LSE and the PV scale, as the producer derives them. + int const r = tid - kComputeThreads; + float const* rm = redx_m + r * kWpb; + float const* rl = redx_l + r * kWpb; + float m = kNegInf; + float l = 0.0F; +#pragma unroll + for (int w4 = 0; w4 < kWpb; w4 += 4) + { + float4 const vm = *reinterpret_cast(rm + w4); + float4 const vl = *reinterpret_cast(rl + w4); + m = fmaxf(m, fmaxf(fmaxf(vm.x, vm.y), fmaxf(vm.z, vm.w))); + l += (vl.x + vl.y) + (vl.z + vl.w); + } + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[r] = scale * vgs; + row_lse[r] = lse; + } + __syncthreads(); // O partials and row statistics visible to the push + } + else + { + if constexpr (kRs2) + { + // ---------------- Split softmax: warps 0..3 rows 0..7, warps 4..7 rows 8..15 (thread = block token) + // -------- + if (warp < 8) + { + int const rg = warp / 4; + int const wq = warp % 4; + int const tk = wq * 32 + lane; + auto rg_bar = [&]() + { + if (rg == 0) + { + asm volatile("bar.sync 10, 128;\n" ::: "memory"); + } + else + { + asm volatile("bar.sync 11, 128;\n" ::: "memory"); + } + }; + int const rs8 = ((lane >> 4) & 1) * 4 + ((lane >> 3) & 1) * 2 + ((lane >> 2) & 1); + mbar_wait_cta(mma_bar, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF(7); + float v[8]; + tmem_ld8(tmem + ((wq * 32) << 16) + 8 * rg, v); + bool const valid = cta_tok0 + tk < kv_max && cta_tok0 + tk < kv_len[0]; + float a[8]; +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] = valid ? v[j2] * sl2e : kNegInf; + a[j2] = v[j2]; + } + float const mw = warp_reduce_scatter8(a, lane); + if ((lane & 3) == 0) + { + red_m[(8 * rg + rs8) * 4 + wq] = mw; + } + rg_bar(); + M3FA_PROF(8); + uint8_t* pdst = smem + S::kP + (tk >> 3) * 128 + (tk & 7) * 16; + float mrow = kNegInf; + if (lane < 8) + { + float4 const mq = *reinterpret_cast(red_m + (8 * rg + lane) * 4); + mrow = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + } + float pa[8]; +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + float const m = __shfl_sync(0xffffffffu, mrow, j2); + pa[j2] = m == kNegInf ? 0.0F : ex2(v[j2] - m); + } + if (tid == 0) + { + M3FA_CLK(41); + } + uint32_t pw[4]; +#pragma unroll + for (int j2 = 0; j2 < 4; ++j2) + { + __half2 const hp = __floats2half2_rn(pa[2 * j2], pa[2 * j2 + 1]); + pw[j2] = *reinterpret_cast(&hp); + } + *reinterpret_cast(pdst + rg * 2048) = make_uint4(pw[0], pw[1], pw[2], pw[3]); + M3FA_PROF_T(18, 0); + if (tid == 0) + { + M3FA_CLK(42); + } + M3FA_FENCE_PROXY_ASYNC(); // P -> tensor core reads + if (tid == 0) + { + M3FA_CLK(43); + } + if (warp == 0) + { + asm volatile("bar.sync 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + if (tid == 0) + { + M3FA_CLK(44); + } + } + else + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + if (M3FA_MEGA_WARPMMA ? warp == 0 : tid == 0) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm_u = tmem; + uint32_t const a0 = smem_u32(smem + S::kV16); + uint32_t const b0 = smem_u32(smem + S::kP); + if (M3FA_MEGA_WARPMMA ? elect_one() : true) + { +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + umma_f16(tm_u + BPC * kRows, umma_desc(a0 + ks * 4096, 2048, 128), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 1, 1), ks > 0); + } + umma_commit(mma_bar); + } + if (M3FA_MEGA_WARPMMA) + { + __syncwarp(); + } + } + M3FA_PROF_T(19, 0); + float const lw = warp_reduce_scatter8(pa, lane); + if ((lane & 3) == 0) + { + red_l[(8 * rg + rs8) * 4 + wq] = lw; + } + rg_bar(); + M3FA_PROF_T(20, 0); + if (wq == 0 && lane < 8) + { + int const r = 8 * rg + lane; + float4 const mq = *reinterpret_cast(red_m + r * 4); + float4 const lq = *reinterpret_cast(red_l + r * 4); + float const m = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + float const l = (lq.x + lq.y) + (lq.z + lq.w); + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[r] = scale * vgs; + row_lse[r] = lse; + } + } + __syncthreads(); // row scales and LSEs visible (PV was issued above) + M3FA_PROF(9); + if (warp < 8) + { + mbar_wait_cta(mma_bar, 1); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + } + else + { + // ---------------- Softmax: thread = block token; the (block, 16-row chunk) pairs are spread over the + // four groups of four compute warps (each group covers the 128 TMEM lanes = block tokens) + // ---------------- + if (warp < kComputeWarps && warp / 4 < kPairs) + { + int const g = warp / 4; + int const wq = warp % 4; + int const tk = wq * 32 + lane; + auto group_bar = [&]() + { + if (g == 0) + { + asm volatile("bar.sync 10, 128;\n" ::: "memory"); + } + else if (g == 1) + { + asm volatile("bar.sync 11, 128;\n" ::: "memory"); + } + else if (g == 2) + { + asm volatile("bar.sync 12, 128;\n" ::: "memory"); + } + else + { + asm volatile("bar.sync 13, 128;\n" ::: "memory"); + } + }; + // Row r(lane) this lane holds after the reduce-scatter: bits 4..1 of the lane. + int const rs_row + = ((lane >> 4) & 1) * 8 + ((lane >> 3) & 1) * 4 + ((lane >> 2) & 1) * 2 + ((lane >> 1) & 1); + mbar_wait_cta(mma_bar, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF(7); + float v[kPairsPerGroup][16]; +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + tmem_ld16(tmem + ((wq * 32) << 16) + h * kRows + c * 16, v[k]); + bool const valid + = cta_tok0 + h * kPage + tk < kv_max && cta_tok0 + h * kPage + tk < kv_len[c]; + float a[16]; +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[k][j] = valid ? v[k][j] * sl2e : kNegInf; + a[j] = v[k][j]; + } + float const mw = warp_reduce_scatter16(a, lane); + if ((lane & 1) == 0) + { + red_m[((h * QG + c) * 16 + rs_row) * 4 + wq] = mw; + } + } + } + group_bar(); + M3FA_PROF(8); + float pa[kPairsPerGroup][16]; +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + uint8_t* pdst = smem + S::kP + h * kRows * kPage * 2 + (tk >> 3) * 128 + (tk & 7) * 16; + // Row maxima: lane j (< 16) reduces row j's four warp partials, then broadcasts it. + float mrow = kNegInf; + if (lane < 16) + { + float4 const mq + = *reinterpret_cast(red_m + ((h * QG + c) * 16 + lane) * 4); + mrow = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + } +#pragma unroll + for (int j = 0; j < 16; ++j) + { + float const m = __shfl_sync(0xffffffffu, mrow, j); + pa[k][j] = m == kNegInf ? 0.0F : ex2(v[k][j] - m); + } + if (tid == 0) + { + M3FA_CLK(41); + } + uint32_t pw[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) + { + __half2 const hp = __floats2half2_rn(pa[k][2 * j], pa[k][2 * j + 1]); + pw[j] = *reinterpret_cast(&hp); + } + // P^T as the MN-major B operand: 8 query rows of one token are one 16-byte core-matrix row. + *reinterpret_cast(pdst + (2 * c) * 2048) = make_uint4(pw[0], pw[1], pw[2], pw[3]); + *reinterpret_cast(pdst + (2 * c + 1) * 2048) + = make_uint4(pw[4], pw[5], pw[6], pw[7]); + M3FA_PROF_T(18, 0); + } + } + if (tid == 0) + { + M3FA_CLK(42); + } + M3FA_FENCE_PROXY_ASYNC(); // P -> tensor core reads + if (tid == 0) + { + M3FA_CLK(43); + } + // PV starts as soon as P and the V tiles are in: thread 0 waits for the softmax warps and the + // producer warps (V); the row sums and statistics below overlap the MMAs. + if (warp == 0) + { + asm volatile("bar.sync 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + else + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + if (tid == 0) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + uint32_t const a0 = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const b0 = smem_u32(smem + S::kP + h * kRows * kPage * 2); +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + umma_f16(tmem + (BPC + h) * kRows, umma_desc(a0 + ks * 4096, 2048, 128), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 1, 1), ks > 0); + } + } + umma_commit(mma_bar); + } + M3FA_PROF_T(19, 0); +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + float const lw = warp_reduce_scatter16(pa[k], lane); + if ((lane & 1) == 0) + { + red_l[((h * QG + c) * 16 + rs_row) * 4 + wq] = lw; + } + } + } + group_bar(); + M3FA_PROF_T(20, 0); + if (wq == 0 && lane < 16) + { +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + int const r = (h * QG + c) * 16 + lane; + float4 const mq = *reinterpret_cast(red_m + r * 4); + float4 const lq = *reinterpret_cast(red_l + r * 4); + float const m = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + float const l = (lq.x + lq.y) + (lq.z + lq.w); + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[h * kRows + c * 16 + lane] = scale * vgs; + row_lse[h * kRows + c * 16 + lane] = lse; + } + } + } + } + __syncthreads(); // row scales and LSEs visible (PV was issued above) + M3FA_PROF(9); + if (warp < kComputeWarps && warp / 4 < kPairs) + { + mbar_wait_cta(mma_bar, 1); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + } + } + M3FA_PROF(3); + } + else + { + asm volatile("cp.async.wait_all;\n" ::: "memory"); +#if M3FA_TRIGGER_AT == 2 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (!FQKV || M3FA_MEGA_TRIGGER == 1) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + for (int r = tid; r < BPC * kRows; r += kThreads) + { + row_lse[r] = kNegInf; + row_scale[r] = 0.0F; + } + __syncthreads(); + } + + // ---------------- Push every partial row to the CTA that merges its head (CTA r merges heads r*BPC..) -------- + if constexpr (S::kRecvAlias) + { + // Peers' partials overlay this CTA's tiles: every CTA's PV MMAs have completed. + asm volatile("barrier.cluster.arrive.aligned;\n" ::: "memory"); + } + if (!(FQKV && M3FA_MEGA_EARLYCLW && !S::kRecvAlias)) + { + asm volatile("barrier.cluster.wait.aligned;\n" ::: "memory"); // peers' mbarriers are initialized + } + M3FA_PROF(4); + uint32_t const recv_base = smem_u32(recv); + if (kRs2 && warp < 8) + { + // Split push: warp group rg pushes rows 8 rg .. 8 rg + 7 (row j goes to CTA j, which merges head j). + int const rg = warp / 4; + int const d = (warp % 4) * 32 + lane; + float v[8]; + if (any_valid) + { + tmem_ld8(tmem + ((warp % 4) * 32 << 16) + BPC * kRows + 8 * rg, v); + } + else + { +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] = 0.0F; + } + } +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] *= row_scale[8 * rg + j2]; + } +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + uint32_t const dst_cta = 8 * rg + j2; + uint32_t const row_addr = recv_base + static_cast((cr * kOPitch + d) * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr, dst_cta)), + "r"(__float_as_uint(v[j2])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + else if (!kRs2 && warp < kComputeWarps && warp / 4 < kPairs) + { + int const g = warp / 4; + int const d = (warp % 4) * 32 + lane; // output dim +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + int const chunk = cr * BPC + h; + float v[16]; + if (any_valid && kRegCore) + { + float const* obuf + = reinterpret_cast(smem + S::kK16 + h * S::kTileBytes + S::kObufOff); +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] = obuf[(c * 16 + j) * kOPitch + d]; + } + } + else if (any_valid) + { + tmem_ld16(tmem + ((warp % 4) * 32 << 16) + (BPC + h) * kRows + c * 16, v); + } + else + { +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] = 0.0F; + } + } +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] *= row_scale[h * kRows + c * 16 + j]; + } + // One scalar store per row: a vector store's lanes would address several destination CTAs at once, + // which measured several times slower. +#pragma unroll + for (int j = 0; j < 16; ++j) + { + uint32_t const dst_cta = j / BPC; + uint32_t const row_addr + = recv_base + static_cast((((chunk * QG + c) * BPC + j % BPC) * kOPitch + d) * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr, dst_cta)), + "r"(__float_as_uint(v[j])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + } + } + else if (tid >= kComputeThreads) + { + // Every partial row's LSE (BPC * kRows of them; more than the producer threads for the widest geometries). + for (int r = tid - kComputeThreads; r < BPC * kRows; r += kThreads - kComputeThreads) + { + int const h = r / kRows; + int const c = (r % kRows) / 16; + int const j = r % 16; + uint32_t const dst_cta = j / BPC; + uint32_t const row_addr + = recv_base + static_cast((((cr * BPC + h) * QG + c) * BPC + j % BPC) * kOPitch * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr + kHeadDim * 4, dst_cta)), + "r"(__float_as_uint(row_lse[r])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + M3FA_PROF(10); +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + if (wr_kind[k] == 1) + { + *reinterpret_cast(wr_ptr[k]) = static_cast(wr_val[k] & 0xFFFFu); + if (lane % 4 == 0) + { + *wr_sptr[k] = static_cast(wr_val[k] >> 16); + } + } + else if (wr_kind[k] == 2) + { + *reinterpret_cast(wr_ptr[k]) = wr_val[k]; + } + } + { + uint32_t ready = 0; +#ifdef M3FA_DBG_HANG + long long spins = 0; +#endif + while (!ready) + { +#ifdef M3FA_DBG_HANG + if (++spins == (1LL << 24)) + { + printf("M3FA hang: block %d thread %d merge recv\n", blockIdx.x, threadIdx.x); + } +#endif + 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}\n" + : "=r"(ready) + : "r"(mbar) + : "memory"); + } + } + M3FA_PROF(5); + + // ---------------- Merge heads [cr * BPC, cr * BPC + BPC) (split-K merge kernel's LSE weighting) ---------------- + constexpr int kMergeRows = QG * BPC; // (token, local head) pairs, 16 chunk weights each + // Wide geometries: chunk weights once per merged row (into the softmax scratch, free by now), then the weighted + // sums. Narrow ones recompute the weights per output element: no extra barrier on the critical path. + constexpr bool kRowWeights = kMergeRows >= 4; + float* wmerge = red_m; // [kMergeRows][kBlocks]; red_m holds BPC * kRows * 4 = kMergeRows * 64 floats + [[maybe_unused]] auto chunk_weights = [&](int mr, float(&lc)[kBlocks]) + { + float lmax = kNegInf; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + lc[c] = recv[(c * kMergeRows + mr) * kOPitch + kHeadDim]; + lmax = fmaxf(lmax, lc[c]); + } + lmax = lmax == kNegInf ? 0.0F : lmax; + float denom = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + lc[c] = ex2(lc[c] - lmax); // now the chunk weight; 0 for an empty chunk + denom += lc[c]; + } + denom = denom > 0.0F ? denom : 1.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + // -0.0 marks an empty chunk (skipped below, as the merge kernel skips zero weights). + lc[c] = lc[c] != 0.0F ? __fdividef(lc[c], denom) : -0.0F; + } + }; + if constexpr (kRowWeights) + { + for (int mr = tid; mr < kMergeRows; mr += kThreads) + { + float lc[kBlocks]; + chunk_weights(mr, lc); +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wmerge[mr * kBlocks + c] = lc[c]; + } + } + __syncthreads(); + } +#pragma unroll + for (int idx = tid; idx < kMergeRows * kHeadDim; idx += kThreads) + { + int const mr = idx >> 7; + int const d = idx & (kHeadDim - 1); + float wn[kBlocks]; + if constexpr (kRowWeights) + { +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wn[c] = wmerge[mr * kBlocks + c]; + } + } + else + { +#ifndef M3FA_MEGA_MERGE_LOCAL +#define M3FA_MEGA_MERGE_LOCAL \ + 0 // 1 (measured +0.6 us slower): every merging thread derives the 16 chunk weights itself (broadcast shared loads, + // independent exps): no shuffle chain; the same arithmetic as chunk_weights +#endif + if constexpr (M3FA_MEGA_MERGE_LOCAL != 0) + { + chunk_weights(mr, wn); + } + else + { + // Narrow geometries: the warp (32 dims of one merged row) derives the 16 chunk weights together, lane c + // holding chunk c's, and broadcasts them. + float lc = (lane < kBlocks) ? recv[(lane * kMergeRows + mr) * kOPitch + kHeadDim] : kNegInf; + float lmax = lc; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + lmax = fmaxf(lmax, __shfl_xor_sync(0xffffffffu, lmax, o)); + } + lmax = lmax == kNegInf ? 0.0F : lmax; + float const wc = lane < kBlocks ? ex2(lc - lmax) : 0.0F; + // The denominator in chunk order from 0.0, as chunk_weights (and the merge kernel) add it: a butterfly + // reduction rounds differently in the last bit. + float denom = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + denom += __shfl_sync(0xffffffffu, wc, c); + } + denom = denom > 0.0F ? denom : 1.0F; + float const wl = wc != 0.0F ? __fdividef(wc, denom) : -0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wn[c] = __shfl_sync(0xffffffffu, wl, c); + } + } + } + if (tid == 0) + { + M3FA_CLK(45); + } + float acc = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + float const prod = recv[(c * kMergeRows + mr) * kOPitch + d] * wn[c]; + acc = __float_as_uint(wn[c]) != 0x80000000u ? acc + prod : acc; + } + if (tid == 0) + { + M3FA_CLK(46); + } + int const tl = mr / BPC; + int const head = cr * BPC + mr % BPC; + __nv_bfloat16 const ob = __float2bfloat16_rn(acc); + int const out_row = tok0 + tfirst + tl; + p.out[static_cast(out_row) * p.out_stride + head * kHeadDim + d] = ob; + if constexpr (MXFP8) + { + float const x = __bfloat162float(ob); + float amax = fabsf(x); +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, mask)); + } + float const sfv = amax * rcp_approx_ftz(448.0F); + int const bits = __float_as_int(sfv); + int ex = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0 ? 1 : 0); + ex = min(ex, 254); + float scale = __int_as_float((254 - ex) << 23); + scale = amax != 0.0F ? scale : 0.0F; + __nv_fp8_e4m3 const q8v(x * scale); + int64_t const q8off = static_cast(out_row) * (p.num_sf_k_tiles * 128) + head * kHeadDim + d; + if constexpr (FQKV) + { + // Megakernel hand-off to the o_proj units (Lamport): this call's parity buffer, the other one reset to + // 0xFF. 0xFF is never written: a NaN code is canonicalized to 0x7F; the UE8M0 code is at most 254. + uint8_t const code = q8v.__x == 0xFFu ? static_cast(0x7Fu) : q8v.__x; + int64_t const cur = (mega_calls & 1u) * kMaxTok * kOpK; + int64_t const nxt = ((mega_calls & 1u) ^ 1u) * kMaxTok * kOpK; + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.q8 + cur + q8off), "h"(static_cast(code)) + : "memory"); + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.q8 + nxt + q8off), "h"(static_cast(0xFFu)) + : "memory"); + } + else + { + p.q8[q8off] = q8v.__x; + } + if (lane == 0) + { + int const k_sf = head * (kHeadDim / 32) + (d >> 5); + int64_t const off = static_cast(out_row / 128) * p.num_sf_k_tiles * 512 + (k_sf / 4) * 512 + + (out_row % 32) * 16 + ((out_row % 128) / 32) * 4 + k_sf % 4; + if constexpr (FQKV) + { + int64_t const cur = (mega_calls & 1u) * kOpSfBytes; + int64_t const nxt = ((mega_calls & 1u) ^ 1u) * kOpSfBytes; + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.sf + cur + off), "h"(static_cast(ex)) + : "memory"); + asm volatile("st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.sf + nxt + off), + "h"(static_cast(0xFFu)) + : "memory"); + } + else + { + p.sf[off] = static_cast(ex); + } + } + } + } + M3FA_PROF(60); + if (any_valid) + { + // Every TMEM read completed above (tcgen05.wait::ld); release the columns. + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + __syncthreads(); + if (FQKV && tid == 0) + { + p.flags[32 + blockIdx.x] = mega_calls + 1u; + } + M3FA_PROF(61); + if (warp == kWarps - 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" ::"r"(*tmem_slot), "r"(S::kTmemCols) + : "memory"); + } + M3FA_PROF(6); +#else + __trap(); +#endif // MINIMAX_M3_ATTN_DEVICE_ENABLED +} + +template +cudaError_t launch_impl(Params const& prm, int num_reqs, bool pdl, cudaStream_t stream, int* max_clusters) +{ + auto kernel = m3_fused_decode_attn_kernel; + constexpr int smem = SmemX::kBytes; + constexpr int kCs = kBlocks / BPC; + static bool const configured = [&]() + { + bool ok = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem) == cudaSuccess; + ok = ok + && (kCs <= 8 + || cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) == cudaSuccess); + if (!ok) + { + cudaGetLastError(); + } + return ok; + }(); + if (!configured) + { + return cudaErrorInvalidConfiguration; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(FQKV ? kQkvGrid : kCs * (QLEN / QG) * num_reqs); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cudaLaunchAttribute attr[2]; + attr[0].id = cudaLaunchAttributeClusterDimension; + attr[0].val.clusterDim.x = kCs; + attr[0].val.clusterDim.y = 1; + attr[0].val.clusterDim.z = 1; + attr[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[1].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + if (max_clusters != nullptr) + { + cfg.numAttrs = 1; + cudaError_t const err = cudaOccupancyMaxActiveClusters(max_clusters, kernel, &cfg); + if (err != cudaSuccess) + { + cudaGetLastError(); + } + return err; + } + cfg.numAttrs = 2; + return cudaLaunchKernelEx(&cfg, kernel, prm); +} + +// The one geometry the target launches: one request, one query token, one 16-CTA attention cluster. +cudaError_t dispatch(Params const& prm, int qlen, int qg, int bpc, int num_reqs, bool pdl, cudaStream_t stream) +{ + if (qlen == 1 && qg == 1 && bpc == 1) + { + return launch_impl<1, 1, 1, true, true>(prm, num_reqs, pdl, stream, nullptr); + } + return cudaErrorInvalidValue; +} + +static_assert(kQkvGrid == minimax_m3_attn_mega::kGrid && kMaxTok == minimax_m3_attn_mega::kMaxTokens + && kQkvN == minimax_m3_attn_mega::kQkvN && kQkvK == minimax_m3_attn_mega::kQkvK + && kOpN == minimax_m3_attn_mega::kOpN && kOpK == minimax_m3_attn_mega::kOpK + && kOpSfBytes == minimax_m3_attn_mega::kOpSfBytes, + "header dimensions"); + +} // namespace + +namespace +{ +// TMA descriptor of a [n, k] E4M3 weight (128 x 128 boxes, 128-byte swizzle), cached per weight address. +bool weight_map(CUtensorMap& map, void const* w, int k, int n) +{ + static std::mutex mu; + static std::unordered_map cache; + std::lock_guard lock(mu); + auto it = cache.find(w); + if (it != cache.end()) + { + map = it->second; + return true; + } + static PFN_cuTensorMapEncodeTiled_v12000 encode = []() + { + void* fn = nullptr; + cudaDriverEntryPointQueryResult q; + cudaGetDriverEntryPointByVersion("cuTensorMapEncodeTiled", &fn, 12000, cudaEnableDefault, &q); + return reinterpret_cast(fn); + }(); + if (encode == nullptr) + { + return false; + } + cuuint64_t dims[2] = {static_cast(k), static_cast(n)}; + cuuint64_t strides[1] = {static_cast(k)}; + cuuint32_t box[2] = {128, 128}; + cuuint32_t elem[2] = {1, 1}; + CUresult const r = encode(&map, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(w), dims, strides, box, elem, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_L2_256B, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + if (r != CUDA_SUCCESS) + { + return false; + } + cache.emplace(w, map); + return true; +} +} // namespace + +cudaError_t invokeMinimaxM3AttnMega(MinimaxM3AttnMegaParams& params, void const* qkvWeight, void const* oWeight, + int qlen, int qg, int bpc, int numReqs, bool pdl, cudaStream_t stream) +{ + if (!weight_map(params.wmap, qkvWeight, kQkvK, kQkvN) || !weight_map(params.wmap_o, oWeight, kOpK, kOpN)) + { + return cudaErrorInvalidValue; + } + params.packed = params.qkv_out; + params.packed_stride = kQkvN; + params.attn_clusters = numReqs * (qlen / qg); + return dispatch(params, qlen, qg, bpc, numReqs, pdl, stream); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/attnMega.h b/cpp/tensorrt_llm/kernels/minimaxM3/attnMega.h new file mode 100644 index 000000000000..d8e267d6e3db --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/attnMega.h @@ -0,0 +1,112 @@ +/* + * 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 minimax_m3_attn_mega +{ +constexpr int kGrid = 112; // CTAs of the persistent grid (7 clusters of 16) +constexpr int kMaxTokens = 8; // rows of each call-parity buffer +constexpr int kQkvN = 2560; // packed QKV(+index) width +constexpr int kQkvK = 6144; // hidden size +constexpr int kOpN = 6144; // o_proj output width +constexpr int kOpK = 2048; // o_proj input width (16 heads x 128) +constexpr int kOpSfBytes = 16 * 512; // one parity of the MXFP8 hand-off's swizzled scales +} // namespace minimax_m3_attn_mega + +//! Kernel parameters of the attention megakernel (the device code's Params). +struct MinimaxM3AttnMegaParams +{ + __nv_bfloat16 const* packed; + int64_t packed_stride; + uint8_t* kv_data; + int64_t d_page, d_plane; + uint8_t* kv_scale; + int64_t s_page, s_plane; + uint8_t* idx_k; + int64_t i_page, i_tok; + int64_t num_pages; + int const* out_cache_loc; + int const* position_ids; + int const* block_table; + int64_t bt_stride; + int bt_cols; + int const* seq_lens; + float const* kv_quant_scale; + float const* k_gs; + float const* v_gs; + __nv_bfloat16 const* q_w; + __nv_bfloat16 const* k_w; + __nv_bfloat16 const* ik_w; + float const* rope; + float eps; + float sm_scale; + __nv_bfloat16* out; + int64_t out_stride; + uint8_t* q8; + uint8_t* sf; + int num_sf_k_tiles; + unsigned long long* prof; // optional per-CTA clock64 probes (-DM3FA_PROFILE builds only) + // Fused QKV projection (FQKV kernels only): out[M, 2560] = mxfp8(x)[M, 6144] @ W^T. + __nv_bfloat16 const* x; // [M, 6144] layer input (the all-reduce output) + int ldx; + int num_tokens; // M + uint8_t const* wsf; // UE8M0 weight scales, 128x4-swizzled + float* part; // [2 parities][splits][8 tokens][2560] partial sums, sentinel at rest + int* counters; // [0]: call parity, [1]: CTAs done with this call (zero at rest) + int attn_clusters; // leading clusters that run the attention + uint8_t const* xq; // mega: the layer input's MXFP8 values [M, 6144] (row stride ldxq) ... + uint8_t const* xsf_act; // ... and its 128x4-swizzled UE8M0 scale atoms (atom k at byte 512 k) + int ldxq; + __nv_bfloat16* qkv_out; // mega: the packed QKV rows [2 parities][8][2560] the attention reads + uint8_t const* wsf_o; // mega: the o_proj weight's scale atoms + __nv_bfloat16* y_out; // mega: the o_proj output [M, 6144] + int* ar_ws; // mega, ar_mode: the next m3_mc_allreduce's multicast exchange (int32 words)... + int const* ar_flags; // ... its call flags (parity of a (token, column slice) slot = flags[t * 8 + slice] & 1) + int ar_rank; + int ar_world; + int ar_maxt; + int ar_mode; // 0: none; 2: push the o_proj rows into the exchange (as m3_mx_gemv's ar mode 2) + unsigned int* flags; // mega: [0, 20) tile counters (4 per call), [32 + CTA] the attention CTAs' targets + long long pf_a0, pf_b0, pf_a1, pf_b1; // mega: L2 prefetch regions issued in the attention window (bytes 0: none) + int pf_hint; // 1: evict_last policy on the prefetch + alignas(64) CUtensorMap wmap; // the [2560, 6144] E4M3 weight, 128 x 128 boxes, 128-byte swizzle + alignas(64) CUtensorMap wmap_o; // mega: the [6144, 2048] E4M3 o_proj weight, same boxes +}; + +//! Launch the megakernel on `stream`. The caller fills `params` except `wmap`, `wmap_o`, `packed`, +//! `packed_stride` and `attn_clusters`, which this derives (the TMA descriptors of the two weights, cached per +//! address). Only (qlen, qg, bpc) = (1, 1, 1) with the MXFP8 hand-off is compiled; anything else, or a failed TMA +//! descriptor encode, returns cudaErrorInvalidValue. +cudaError_t invokeMinimaxM3AttnMega(MinimaxM3AttnMegaParams& params, void const* qkvWeight, void const* oWeight, + int qlen, int qg, int bpc, int numReqs, bool pdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/draftStep.cu b/cpp/tensorrt_llm/kernels/minimaxM3/draftStep.cu new file mode 100644 index 000000000000..4df471b7a78f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/draftStep.cu @@ -0,0 +1,576 @@ +/* + * 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/minimaxM3/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 minimaxM3 +{ +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); +} + +// 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][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 + 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 const v = 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; + } + } +#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 minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/draftStep.h b/cpp/tensorrt_llm/kernels/minimaxM3/draftStep.h new file mode 100644 index 000000000000..d6f4c1dd67d9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/draftStep.h @@ -0,0 +1,99 @@ +/* + * 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 minimaxM3 +{ +//! 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; +}; + +//! 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 minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.cu b/cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.cu new file mode 100644 index 000000000000..4383bb0fb2e6 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.cu @@ -0,0 +1,3679 @@ +/* + * 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. + */ + +// Fused short-sequence decode attention of the MiniMax-M3 MSA sparse layers (packed QKV input). + +#include "tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +// The kernel uses tcgen05/TMEM (including block-scaled MMAs, which family targets lack), cluster DSMEM and E2M1 +// conversions: its body is compiled for the arch-specific sm_100a / sm_103a targets only (see the minimax_m3_attn_src +// library in kernels/CMakeLists.txt). Any other target compiles a kernel that traps, so a launch on a build without +// those targets fails loudly instead of returning untouched outputs. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030) \ + && (defined(__CUDA_ARCH_SPECIFIC__) || defined(__CUDA_ARCH_FEAT_SM100_ALL) || defined(__CUDA_ARCH_FEAT_SM103_ALL)) +#define MINIMAX_M3_ATTN_DEVICE_ENABLED 1 +#else +#define MINIMAX_M3_ATTN_DEVICE_ENABLED 0 +#endif + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr int kHeadDim = 128; +constexpr int kRotDim = 64; +constexpr int kPage = 128; +constexpr int kHq = 16; +constexpr int kColK = kHq; // packed head slots: [Q x16 | K | V | idxQ | idxK] +constexpr int kColV = kHq + 1; +constexpr int kColIdxK = kHq + 3; +constexpr int kBlocks = 16; // KV blocks per request == topk == cluster size +constexpr int kComputeWarps = 16; // QK / softmax / PV and the dequant (one 16-byte chunk per thread) +constexpr int kWarps = kComputeWarps + 4; // plus producer-only warps, so new K/V/index-K never queue behind a Q head +constexpr int kThreads = kWarps * 32; +constexpr int kComputeThreads = kComputeWarps * 32; +constexpr int kOPitch = 132; // floats per partial-O row +constexpr float kNegInf = -__builtin_huge_valf(); + +using Params = MinimaxM3FusedDecodeAttnParams; + +// Optional per-CTA clock64 probes for development (-DM3FA_PROFILE and a prof tensor): [CTA][64] timestamps. +// -DM3FA_PROFILE_GT records %globaltimer (ns, comparable across SMs and kernels) instead of clock64. +#ifdef M3FA_PROFILE_GT +#define M3FA_PROFILE +#define M3FA_NOW(v) asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(v)) +#else +#define M3FA_NOW(v) v = clock64() +#endif +#ifdef M3FA_PROFILE +#define M3FA_PROF_T(i, thread) \ + do \ + { \ + if (p.prof != nullptr && threadIdx.x == (thread)) \ + { \ + unsigned long long now_; \ + M3FA_NOW(now_); \ + p.prof[blockIdx.x * 64 + (i)] = now_; \ + } \ + } while (0) +#else +#define M3FA_PROF_T(i, thread) \ + do \ + { \ + } while (0) +#endif +#define M3FA_PROF(i) M3FA_PROF_T(i, 0) +// clock64 probes (development): slot i of this CTA, written by every calling thread. +#ifdef M3FA_PROFILE +#define M3FA_CLK(i) \ + do \ + { \ + if (p.prof != nullptr) \ + { \ + p.prof[blockIdx.x * 64 + (i)] = clock64(); \ + } \ + } while (0) +#else +#define M3FA_CLK(i) \ + do \ + { \ + } while (0) +#endif + +#ifndef M3FA_MEGA_EARLYCLW +#define M3FA_MEGA_EARLYCLW 1 // pass the cluster barrier (peers' merge mbarriers initialized) early, not before the push +#endif +#ifndef M3FA_V15_EARLYV +#define M3FA_V15_EARLYV 0 // packed path: dequantize V before the PDL wait (switch; off by default) +#endif +#ifndef M3FA_MEGA_TRIGGER +// Megakernel PDL release: 0 o_proj CTAs after their activation (attention CTAs at exit); 1 as 0 with the attention +// CTAs after their KV reads; 2 every CTA right after its griddepcontrol.wait; 3 o_proj CTAs after their epilogue. +#define M3FA_MEGA_TRIGGER 3 +#endif +#ifndef M3FA_TRIGGER_AT +#define M3FA_TRIGGER_AT 2 // where the PDL dependents are released: 0 entry, 1 after the wait, 2 after the KV reads +#endif + +// Generic-proxy shared-memory writes (fp16 tiles, Q, P) become visible to the tensor core (async proxy). +#define M3FA_FENCE_PROXY_ASYNC() asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory") + +__device__ __forceinline__ void cp_async16(void* dst, void const* src) +{ + uint32_t const d = static_cast(__cvta_generic_to_shared(dst)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(d), "l"(src) : "memory"); +} + +__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__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +// One packed E2M1 byte -> half2 (low nibble in the low half). +__device__ __forceinline__ uint32_t e2m1x2_to_h2(uint32_t byte) +{ + uint32_t out; + asm("{\n .reg .b8 src;\n mov.b32 {src, _, _, _}, %1;\n cvt.rn.f16x2.e2m1x2 %0, src;\n}\n" : "=r"(out) : "r"(byte)); + return out; +} + +__device__ __forceinline__ uint32_t hmul2_u(uint32_t a, uint32_t b) +{ + __half2 r = __hmul2(*reinterpret_cast<__half2*>(&a), *reinterpret_cast<__half2*>(&b)); + return *reinterpret_cast(&r); +} + +__device__ __forceinline__ uint32_t e4m3_to_h2_splat(uint32_t byte) +{ + __half_raw const h = __nv_cvt_fp8_to_halfraw(static_cast<__nv_fp8_storage_t>(byte), __NV_E4M3); + return static_cast(h.x) | (static_cast(h.x) << 16); +} + +__device__ __forceinline__ uint32_t e4m3x2_to_h2(uint32_t two_bytes) +{ + __half2_raw const h = __nv_cvt_fp8x2_to_halfraw2(static_cast<__nv_fp8x2_storage_t>(two_bytes), __NV_E4M3); + return static_cast(h.x) | (static_cast(h.y) << 16); +} + +// ---- Exact copies of the C++ producer's helpers (fusedQKNormRopeKernel.cu / quantization.cuh) ---- +template +__device__ __forceinline__ uint32_t cvt_sf_group_shfl_mask() +{ + constexpr uint32_t groupSize = static_cast(NUM_THREADS_PER_SF); + constexpr uint32_t groupMask = 0xFFFFFFFFU >> (32U - groupSize); + uint32_t laneId = 0; + asm("mov.u32 %0, %%laneid;" : "=r"(laneId)); + return groupMask << (laneId & ~(groupSize - 1U)); +} + +__device__ __forceinline__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) +{ + 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"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]), "f"(array[4]), "f"(array[5]), "f"(array[6]), + "f"(array[7])); + return val; +} + +__device__ __forceinline__ uint16_t quantizeMinimaxM3Fp4x4( + float const (&elements)[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(elements[i]); + localMax = fmaxf(localMax, fabsf(__bfloat162float(rounded[i]))); + } + uint32_t const groupMask = cvt_sf_group_shfl_mask<4>(); + 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 quantValues[8] = {}; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + quantValues[i] = __bfloat162float(rounded[i]) * outputScale; + } + return static_cast(fp32_vec_to_e2m1(quantValues)); +} + +// RMSNorm (gemma 1+w) + NeoX RoPE over the first 64 dims, exactly as the producer: lane owns 4 dims. +__device__ __forceinline__ void norm_rope( + float (&e)[4], float const (&w)[4], float eps, float const (&cs)[4], float const (&sn)[4], int lane) +{ + float ss = 0.0F; + ss += e[0] * e[0] + e[1] * e[1]; + ss += e[2] * e[2] + e[3] * e[3]; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + ss = ss + __shfl_xor_sync(0xffffffffu, ss, mask, 32); + } + float const rmsReciprocal = rsqrtf(ss / static_cast(kHeadDim) + eps); +#pragma unroll + for (int i = 0; i < 4; ++i) + { + e[i] *= rmsReciprocal * (1.0F + w[i]); + } + __syncwarp(); + constexpr int kPairOffset = (kRotDim / 2) / 4; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + int const dim = lane * 4 + i; + float paired = __shfl_xor_sync(0xffffffffu, e[i], kPairOffset); + if (dim < kRotDim) + { + bool const firstHalf = dim < kRotDim / 2; + if (firstHalf) + { + paired = -paired; + } + e[i] = e[i] * cs[i] + paired * sn[i]; + } + } + __syncwarp(); +} + +__device__ __forceinline__ void load_w4(float (&w)[4], __nv_bfloat16 const* src, int lane) +{ + uint2 const raw = *reinterpret_cast(src + lane * 4); + float2 const a = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw.x)); + float2 const b = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw.y)); + w[0] = a.x; + w[1] = a.y; + w[2] = b.x; + w[3] = b.y; +} + +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ uint32_t mapa_shared(uint32_t addr, uint32_t rank) +{ + uint32_t r; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n" : "=r"(r) : "r"(addr), "r"(rank)); + return r; +} + +// QLEN query tokens per request; QG of them per cluster (QLEN / QG clusters per request); BPC KV blocks per CTA +// (cluster of kBlocks / BPC CTAs). +// Byte offset of element (r, k) in a UMMA no-swizzle K-major operand tile of 128 k-elements per row: 8x8 core +// matrices of 16-byte rows, 128 bytes apart along k (LBO) and 2048 bytes apart per 8 rows (SBO). The same formula +// with r = token, k = dim stores V so that it is the MN-major (dims contiguous) A operand of O^T = V^T P^T. +__device__ __forceinline__ int kmaj_off(int r, int k) +{ + return (r >> 3) * 2048 + (k >> 3) * 128 + (r & 7) * 16 + (k & 7) * 2; +} + +__device__ __forceinline__ uint64_t umma_desc(uint32_t addr, uint32_t lbo, uint32_t sbo) +{ + return static_cast((addr >> 4) & 0x3FFFu) | (static_cast((lbo >> 4) & 0x3FFFu) << 16) + | (static_cast((sbo >> 4) & 0x3FFFu) << 32) | (1ull << 46); +} + +// kind::f16 instruction descriptor: F32 accumulate, F16 operands, B K-major, M = 128. +__host__ __device__ constexpr uint32_t umma_idesc(int n, int a_mn_major, int b_mn_major) +{ + return (1u << 4) | (static_cast(a_mn_major) << 15) | (static_cast(b_mn_major) << 16) + | (static_cast(n >> 3) << 17) | (static_cast(128 >> 4) << 24); +} + +// 8-row version for the split softmax: the same lane-bit order (16, 8, 4, then 2 and 1 in full) as +// warp_reduce_scatter16, so every row is combined in the same association (fp add is commutative, not associative). +// After it lanes with bits 4..2 = (b4, b3, b2) hold row 4*b4 + 2*b3 + b2. +template +__device__ __forceinline__ float warp_reduce_scatter8(float (&a)[8], int lane) +{ +#pragma unroll + for (int step = 0; step < 3; ++step) + { + int const mask = 16 >> step; + int const half = 4 >> step; + bool const upper = (lane & mask) != 0; +#pragma unroll + for (int i = 0; i < half; ++i) + { + float const send = upper ? a[i] : a[i + half]; + float const keep = upper ? a[i + half] : a[i]; + float const recv = __shfl_xor_sync(0xffffffffu, send, mask); + a[i] = kMax ? fmaxf(keep, recv) : keep + recv; + } + } + float x = a[0]; + float const o2 = __shfl_xor_sync(0xffffffffu, x, 2); + x = kMax ? fmaxf(x, o2) : x + o2; + float const o1 = __shfl_xor_sync(0xffffffffu, x, 1); + return kMax ? fmaxf(x, o1) : x + o1; +} + +// Reduce 16 per-lane values over the warp, one value per row: after the reduce-scatter lane l holds row +// 8*b4 + 4*b3 + 2*b2 + b1 (b_k = bit k of l) reduced over all 32 lanes. 16 shuffles instead of 80. +template +__device__ __forceinline__ float warp_reduce_scatter16(float (&a)[16], int lane) +{ +#pragma unroll + for (int step = 0; step < 4; ++step) + { + int const mask = 16 >> step; + int const half = 8 >> step; + bool const upper = (lane & mask) != 0; +#pragma unroll + for (int i = 0; i < half; ++i) + { + float const send = upper ? a[i] : a[i + half]; + float const keep = upper ? a[i + half] : a[i]; + float const recv = __shfl_xor_sync(0xffffffffu, send, mask); + a[i] = kMax ? fmaxf(keep, recv) : keep + recv; + } + } + float const other = __shfl_xor_sync(0xffffffffu, a[0], 1); + return kMax ? fmaxf(a[0], other) : a[0] + other; +} + +__device__ __forceinline__ void umma_f16( + uint32_t tmem_d, uint64_t adesc, uint64_t bdesc, uint32_t idesc, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %4, 0;\n tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(idesc), "r"(accumulate) + : "memory"); +} + +__device__ __forceinline__ void umma_commit(uint32_t bar) +{ + asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];\n" ::"r"(bar) : "memory"); +} + +__device__ __forceinline__ void mbar_wait_cta(uint32_t bar, uint32_t parity) +{ + uint32_t ready = 0; +#ifdef M3FA_DBG_HANG + long long spins = 0; +#endif + while (!ready) + { +#ifdef M3FA_DBG_HANG + if (++spins == (1LL << 24)) + { + printf("M3FA hang: block %d thread %d mbar 0x%x parity %u\n", blockIdx.x, threadIdx.x, bar, parity); + } +#endif + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(bar), "r"(parity) + : "memory"); + } +} + +// Spin without suspending (v52's acc_done wait): a suspended try_wait wakes late on a tcgen05.commit arrival. +#ifndef M3FA_EPI_SPIN +#define M3FA_EPI_SPIN 1 +#endif +__device__ __forceinline__ void mbar_wait_epi(uint32_t bar, uint32_t parity) +{ + if (!M3FA_EPI_SPIN) + { + mbar_wait_cta(bar, parity); + return; + } + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(bar), "r"(parity) + : "memory"); + } +} + +__device__ __forceinline__ void tmem_ld8(uint32_t addr, float (&v)[8]) +{ + uint32_t r[8]; + asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]) + : "r"(addr)); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < 8; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +__device__ __forceinline__ void tmem_ld16(uint32_t addr, float (&v)[16]) +{ + uint32_t r[16]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, " + "%15}, [%16];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]) + : "r"(addr)); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < 16; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +// Warp-level fp16 MMA (m16n8k16, fp32 accumulate) and its shared-memory fragment loads. +__device__ __forceinline__ void ldsm_x4(uint32_t addr, uint32_t (&r)[4]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void ldsm_x2(uint32_t addr, uint32_t (&r)[2]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" : "=r"(r[0]), "=r"(r[1]) : "r"(addr)); +} + +__device__ __forceinline__ void ldsm_x2_trans(uint32_t addr, uint32_t (&r)[2]) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0, %1}, [%2];\n" + : "=r"(r[0]), "=r"(r[1]) + : "r"(addr)); +} + +__device__ __forceinline__ void mma_f16_16816(float (&d)[4], uint32_t const (&a)[4], uint32_t const (&b)[2]) +{ + 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};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +// Register attention core (warp MMAs, no TMEM) for up to two tokens per cluster; the tcgen05 core above that. +// Off by default: on GB300 it measured ~0.25 us slower than the tcgen05 core at M=1 (GEMV + attention chain). +#ifndef M3FA_REG_CORE +#define M3FA_REG_CORE 0 +#endif +template +struct RegCore +{ + static constexpr bool kOn = M3FA_REG_CORE != 0 && QG <= 2; +}; + +template +struct Smem +{ + static constexpr int kRows = 16 * QG; // query rows (token, head) of this cluster = UMMA N + static constexpr int kTok = BPC * kPage; // block tokens per CTA + static constexpr int kTileBytes = kPage * kHeadDim * 2; // one fp16 K or V block tile in UMMA layout + static constexpr int kK16 = 0; + static constexpr int kV16 = kK16 + BPC * kTileBytes; + static constexpr int kQ = kV16 + BPC * kTileBytes; // Q rows, UMMA K-major [row][dim] + // P^T per block, UMMA MN-major [token][row], over the block's K tile: P is written after the QK MMAs (which + // read every K tile) have completed. + static constexpr int kP = kK16; + // Packed V tiles as copied from the cache: [block][128 rows x 64 data bytes | 1 KB of 4x4-quad scales]. + static constexpr int kRawV = kQ + kRows * kHeadDim * 2; + static constexpr int kRawVBlock = kPage * 64 + kPage * 8; + static constexpr int kRowScale = kRawV + BPC * kRawVBlock; + static constexpr int kRowLse = kRowScale + BPC * kRows * 4; + // Per-warp softmax partials [block][row][4 warps]: max, then sum. + static constexpr int kRedM = (kRowLse + BPC * kRows * 4 + 15) / 16 * 16; + static constexpr int kRedL = kRedM + BPC * kRows * 16; + // Partials pushed by every CTA of the cluster: [chunk][token][local head][kOPitch] floats, lse at column 128. + // When they do not fit beside the tiles they overlay the K/V tiles (and P, Q): peers then push only once every + // CTA of the cluster has completed its PV MMAs (a cluster barrier right before the push). + // Register core: per-warp row maxima and sums [block][row][warp of the block], two arrays. + static constexpr int kRedX = kRedL + BPC * kRows * 16; + static constexpr int kRedXBytes = RegCore::kOn ? 2 * BPC * kRows * (16 / BPC) * 4 : 0; + static constexpr int kStatsEnd = kRedX + kRedXBytes; + // Register core: P (fp16, [row][token] core matrices) at the start of its block's K tile, the scaled-free O + // partial ([row][kOPitch] fp32) behind it. + static constexpr int kObufOff = 8192; + static_assert(!RegCore::kOn || (kRows * kPage * 2 <= kObufOff && kObufOff + kRows * kOPitch * 4 <= kTileBytes), + "register-core P and O staging fit in a K tile"); + static constexpr int kRecvBytes = kBlocks * QG * BPC * kOPitch * 4; + static constexpr bool kRecvAlias = kStatsEnd + kRecvBytes + 32 > 227 * 1024; + static constexpr int kRecv = kRecvAlias ? 0 : kStatsEnd; + static constexpr int kMbar = kRecvAlias ? kStatsEnd : kStatsEnd + kRecvBytes; // merge, MMA mbarriers, TMEM addr + static constexpr int kBytes = kMbar + 32; + static constexpr int kTmemCols = BPC * 2 * kRows <= 32 ? 32 + : BPC * 2 * kRows <= 64 ? 64 + : BPC * 2 * kRows <= 128 ? 128 + : BPC * 2 * kRows <= 256 ? 256 + : 512; + static_assert(BPC * 2 * kRows <= 512, "TMEM columns"); + static_assert(!kRecvAlias || kRecvBytes <= kRawV, "aliased partials must fit in the tiles, Q and P"); + static_assert(kRows * kPage * 2 <= kTileBytes, "P^T fits in the K tiles"); + static_assert(kBytes <= 227 * 1024, "shared memory budget"); + static_assert(QLEN % QG == 0 && (BPC == 1 || BPC == 2), "unsupported geometry"); +}; + +// The packed NVFP4 bytes of one 16-byte chunk (32 dims of one block row) and its two block scales. +struct RawChunk +{ + uint4 data; + uint32_t scales; // byte 0: dims [0, 16), byte 1: dims [16, 32) of the chunk +}; + +// Chunk q of a block tile: lanes 8j..8j+7 of a warp take 8 consecutive rows of one 32-dim chunk, so every 16-byte +// fp16 store phase writes 8 distinct core-matrix rows (conflict-free in the UMMA layout). +__device__ __forceinline__ int chunk_row(int q) +{ + return (q >> 5) * 8 + (q & 7); +} + +__device__ __forceinline__ int chunk_col(int q) +{ + return (q >> 3) & 3; +} + +__device__ __forceinline__ RawChunk load_chunk( + uint8_t const* data_page, uint8_t const* scale_page, int row, int cc, bool is_v) +{ + RawChunk c; + c.data = *reinterpret_cast(data_page + row * 64 + cc * 16); + if (is_v) + { + // V scales: 4x4 token-quad order [row / 4][scale column][row % 4]. + uint8_t const* sc = scale_page + (row >> 2) * 32 + (row & 3); + c.scales = static_cast(sc[(cc * 2) * 4]) | (static_cast(sc[(cc * 2 + 1) * 4]) << 8); + } + else + { + c.scales = *reinterpret_cast(scale_page + row * 8 + cc * 2); + } + return c; +} + +__device__ __forceinline__ void store_chunk_fp16(uint8_t* tile, int row, int cc, RawChunk const& c, bool valid) +{ + uint32_t const s0 = e4m3_to_h2_splat(c.scales & 0xFFu); + uint32_t const s1 = e4m3_to_h2_splat((c.scales >> 8) & 0xFFu); + uint32_t const wd[4] = {c.data.x, c.data.y, c.data.z, c.data.w}; + uint32_t out[16]; +#pragma unroll + for (int i = 0; i < 16; ++i) + { + uint32_t const b = (wd[i >> 2] >> (8 * (i & 3))) & 0xFFu; + out[i] = valid ? hmul2_u(e2m1x2_to_h2(b), i < 8 ? s0 : s1) : 0u; + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + *reinterpret_cast(tile + kmaj_off(row, cc * 32 + j * 8)) + = make_uint4(out[4 * j], out[4 * j + 1], out[4 * j + 2], out[4 * j + 3]); + } +} + +// ---------------- Fused QKV projection: block-scaled MXFP8 on tcgen05 (N=2560, K=6144) ---------------- +// out[M, 2560] = mxfp8(x)[M, 6144] @ W^T, the computation of the CUTLASS MXFP8 GEMM: W is the A operand (128 output +// rows per MMA), the quantized activation the B operand (8 token columns), FP32 accumulation in TMEM. A unit is one +// (128-row tile, K split) pair; the 20 x kQkvSplits units run on the last CTAs of the grid. Before the PDL wait a +// CTA streams its unit's weight tiles and scale atoms through a shared-memory ring into TMEM; after it, it quantizes +// its x slice (mxfp8_quantize math), runs the MMAs and stores its FP32 partial rows. +// +// Hand-off without a flag: the partial buffers are double buffered by call parity and hold a sentinel (an FP32 NaN +// the MMA never produces) at rest, so a consumer polls the partials themselves. The producers of call n reset the +// other parity's buffers after their PDL wait; call n - 1, the last reader of those, is complete by then. +constexpr int kQkvN = 2560; +constexpr int kQkvK = 6144; +constexpr int kQkvKTiles = kQkvK / 128; +constexpr int kQkvRowTiles = kQkvN / 128; +#ifndef M3FA_QKV_SPLITS +#define M3FA_QKV_SPLITS 4 +#endif +#ifndef M3FA_QKV_RING +#define M3FA_QKV_RING kQkvMaxTiles // every tile of the unit in flight at once +#endif +#ifndef M3FA_QKV_NACC +#define M3FA_QKV_NACC 1 +#endif +#ifndef M3FA_QKV_SCALED_MMA +#define M3FA_QKV_SCALED_MMA 1 // 0: unscaled per-block MMAs with the scales applied in the epilogue (see below) +#endif +#ifndef M3FA_QKV_PIPE +// 1: the weight TMA is issued first at entry; the activation has its own region, copied right after the PDL wait; +// each weight tile's MMAs issue as soon as it lands (v52's order). 0: tiles into TMEM first, activation over the ring. +#define M3FA_QKV_PIPE 1 +#endif +#ifndef M3FA_QKV_A_SMEM +#define M3FA_QKV_A_SMEM 1 // scaled MMAs only; 1: they read the weight tiles from shared memory (no copy into TMEM) +#endif +// Default: block-scaled kind::mxf8f6f4 MMAs reading A (the weight tiles) in place from the ring, issued per tile as +// it lands (M3FA_QKV_PIPE), one accumulator; M3FA_QKV_A_SMEM=0 copies each tile into TMEM first. Alternative +// (M3FA_QKV_SCALED_MMA=0): one unscaled kind::f8f6f4 MMA per 32-wide K block into its own accumulator, A read in place +// from shared memory, the epilogue applying both UE8M0 scales. On GB300 an M128 x N8 x K32 MMA occupies the tensor pipe +// ~105 cycles block-scaled and ~60 unscaled, independent of N up to 64; the unscaled form's epilogue (320 accumulator +// columns) costs more than it saves. +constexpr bool kQkvScaled = M3FA_QKV_SCALED_MMA != 0; +constexpr bool kQkvASmem = !kQkvScaled || M3FA_QKV_A_SMEM != 0; +constexpr int kQkvSplits = M3FA_QKV_SPLITS; +constexpr int kQkvUnits = kQkvRowTiles * kQkvSplits; +constexpr int kQkvMaxTiles = (kQkvKTiles + kQkvSplits - 1) / kQkvSplits; +constexpr int kQkvRing = M3FA_QKV_RING; +constexpr int kQkvNacc = M3FA_QKV_NACC; +constexpr int kQkvGrid = 112; // 7 clusters of 16 (or 14 of 8): one wave with the attention +constexpr int kMaxTok = 8; +constexpr int kQkvWTile = 128 * 128; // E4M3 bytes of one 128 x 128 weight tile +constexpr int kQkvSfAtom = 512; // one 128-row x 4-block UE8M0 scale atom +constexpr uint32_t kQkvSentinel = 0xFFFFFFFFu; +constexpr int64_t kQkvPartStride = static_cast(kQkvSplits) * kMaxTok * kQkvN; // one parity +// TMEM columns of the one 512-column allocation of a CTA: attention [0, 64), then the GEMM's accumulators, scale +// atoms (A per tile, B per tile) and the A tiles (32 columns of 128 x 128 E4M3 each). +constexpr int kTmemAll = 512; +// Megakernel: the attention CTAs run no projection unit, so the GEMM starts at column 0 of a projection CTA. +constexpr int kTAcc = 0; +constexpr int kTSfa = kTAcc + 16; +constexpr int kTSfb = kTSfa + 4 * kQkvMaxTiles; +constexpr int kTA = (kTSfb + 4 * kQkvMaxTiles + 31) / 32 * 32; +static_assert(!kQkvScaled || kTA + 32 * kQkvMaxTiles <= kTmemAll, "TMEM columns"); +static_assert(kQkvScaled || kTAcc + 32 * kQkvMaxTiles <= kTmemAll, "TMEM columns (one accumulator per K block)"); +static_assert(!kQkvASmem || kQkvRing == kQkvMaxTiles, "shared-memory A needs every tile resident"); +static_assert(kQkvUnits <= kQkvGrid, "at most one unit per CTA"); +static_assert(kQkvNacc == 1 || kQkvNacc == 2 || kQkvNacc == 4, "accumulator count"); +static_assert(!M3FA_QKV_PIPE || kQkvScaled, "the pipelined projection uses the block-scaled MMAs"); +constexpr int kQkvOwnRows = 128 / kQkvSplits; // rows of a tile whose split-K sum a CTA owns +// Megakernel o_proj (v52 m3_mx_gemv for [6144, 2048]: split 2, contiguous k-tiles, one accumulator): unit +// o = CTA - 16 (the last 96 CTAs), row tile o / 2, split o % 2. +constexpr int kOpN = 6144; +constexpr int kOpK = 2048; +constexpr int kOpKTiles = kOpK / 128; +constexpr int kOpSplits = 2; +constexpr int kOpTiles = kOpKTiles / kOpSplits; +constexpr int kOpOwnRows = 128 / kOpSplits; +constexpr int kOpFirstCta = kQkvGrid - kOpN / 128 * kOpSplits; // 16 +constexpr int kOpSfBytes = kOpKTiles * 512; // one parity of the MXFP8 hand-off's scales +#ifndef M3FA_OP_L2PF +#define M3FA_OP_L2PF \ + 0 // 1: QKV CTAs prefetch their o_proj weight slice into L2 before the PDL wait (competes + // with the QKV weight stream for HBM behind the MoE) +#endif +// o_proj TMEM columns (after the QKV unit's MMAs completed): accumulator, A / B scale atoms, A (8 tiles x 32) +constexpr int kOAcc = 0; +constexpr int kOSfa = 16; +constexpr int kOSfb = kOSfa + 4 * kOpTiles; +constexpr int kOA = 128; +static_assert(kOA + 32 * kOpTiles <= kTmemAll && kOSfb + 4 * kOpTiles <= kOA, "o_proj TMEM columns"); +constexpr int kQkvMbStride = kMaxTok + 4; // mailbox row pitch (fp32 words) + +__host__ __device__ constexpr int qkv_tile0(int s) +{ + return s * kQkvKTiles / kQkvSplits; +} + +__device__ __forceinline__ int qkv_unit() +{ + int const u = kQkvGrid - 1 - static_cast(blockIdx.x); + return u < kQkvUnits ? u : -1; +} + +// Attention layout plus, for the fused projection, the weight ring in front (the attention tiles reuse that space +// once the ring has drained into TMEM); the activation tiles, the pre-wait V copy and the barriers sit behind. +template +struct SmemX : Smem +{ + using B = Smem; + static constexpr int kRingW = 0; // [ring][128 rows][128 B], 128-byte swizzle + static constexpr int kRingSf = kRingW + kQkvRing * kQkvWTile; // [ring][512 B] + static constexpr int kRingEnd = kRingSf + kQkvRing * kQkvSfAtom; + // The activation tiles overlay the ring: written after the PDL wait, once every weight tile went on into TMEM. + static constexpr int kXt = kRingW; // [tile][8][128 B] + static constexpr int kXsf = kXt + kQkvMaxTiles * 1024; // [tile][512 B] + static constexpr int kXsc = (kRingEnd + 127) / 128 * 128; + // Peers push into recv at any time (an empty CTA right away), so it stays clear of the ring. + static constexpr int kRecv = FQKV ? kXsc : B::kRecv; + static constexpr int kRawV = FQKV ? (kRecv + kBlocks * QG * BPC * kOPitch * 4 + 127) / 128 * 128 : B::kRawV; + // Pipelined projection: the activation tiles and scale atoms in their own region (a projection CTA runs no + // attention, so it overlays the attention's partial-O and V regions), 1024-byte aligned for the swizzle. + static constexpr int kXA = FQKV ? (kRecv + 1023) / 1024 * 1024 : 0; + static constexpr int kXASf = kXA + kQkvMaxTiles * 1024; + static constexpr int kXAEnd = kXASf + kQkvMaxTiles * kQkvSfAtom; + static constexpr int kMbarA = FQKV ? kRawV + BPC * B::kRawVBlock : B::kMbar; + static constexpr int kMbar = FQKV && M3FA_QKV_PIPE && kXAEnd > kMbarA ? (kXAEnd + 127) / 128 * 128 : kMbarA; + static constexpr int kGBar = kMbar + 32; // full[ring], empty[ring], a_ready, acc_done, tile_done[tiles] + static constexpr int kRedBar = kGBar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); // split-K partials in + static constexpr int kOBar = kRedBar + 8; // o_proj: full[8], a_ready, acc_done, reduce + static constexpr int kMailbox = (kOBar + 8 * (kOpTiles + 3) + 15) / 16 * 16; // [split][own row][pitch] fp32 + // o_proj weight ring = ring tiles [0, 8); its activation tiles, scale atoms and mailbox behind it. + static constexpr int kOXt = kRingW + kOpTiles * kQkvWTile; + static constexpr int kOXsf = kOXt + kOpTiles * 1024; + static constexpr int kOMailbox = kOXsf + kOpTiles * kQkvSfAtom; + static_assert(!FQKV || kOMailbox + kOpSplits * kOpOwnRows * kQkvMbStride * 4 <= kRingW + kQkvRing * kQkvWTile, + "o_proj activation and mailbox fit behind its ring"); + static constexpr int kBytes = FQKV ? kMailbox + kQkvSplits * kQkvOwnRows * kQkvMbStride * 4 : B::kBytes; + static constexpr int kTmemCols = FQKV ? kTmemAll : B::kTmemCols; + static constexpr bool kRecvAlias = !FQKV && B::kRecvAlias; + // The epilogue reads the weight scale atoms while the attention already writes its tiles. + static_assert(!FQKV || kQkvScaled || B::kRecv <= kRingSf, "attention regions overlap the scale atoms"); + // The other attention regions overlay the ring: a unit's CTA writes them only once the ring has drained. + static_assert(kBytes <= 227 * 1024, "shared memory budget"); +}; + +__device__ __forceinline__ void mbar_init(uint32_t bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(bar), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbar_expect_tx(uint32_t bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(bar), "r"(bytes) : "memory"); +} + +__device__ __forceinline__ void bulk_g2s(void* dst, void const* src, uint32_t bytes, uint32_t bar) +{ + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], " + "%4;\n" ::"r"(smem_u32(dst)), + "l"(src), "r"(bytes), "r"(bar), "l"(0x12F0000000000000ull) // EVICT_FIRST: weights are read once per layer + : "memory"); +} + +__device__ __forceinline__ void tma_load_2d(uint32_t dst, void const* map, int c0, int c1, uint32_t bar) +{ + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3}], [%4], %5;\n" ::"r"(dst), + "l"(map), "r"(c0), "r"(c1), "r"(bar), "l"(0x12F0000000000000ull) + : "memory"); +} + +// K-major operand in the 128-byte swizzle: rows of 128 bytes, 8-row groups 1024 bytes apart. +__device__ __forceinline__ uint64_t umma_desc_sw128(uint32_t addr) +{ + return umma_desc(addr, 16, 1024) | (2ull << 61); +} + +__device__ __forceinline__ void tmem_cp_128x256b(uint32_t taddr, uint64_t desc) +{ + asm volatile("tcgen05.cp.cta_group::1.128x256b [%0], %1;\n" ::"r"(taddr), "l"(desc)); +} + +__device__ __forceinline__ void tmem_cp_sf(uint32_t taddr, uint32_t saddr) +{ + asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;\n" ::"r"(taddr), "l"(umma_desc(saddr, 16, 128))); +} + +// Block-scaled MXFP8 instruction descriptor: E4M3 x E4M3, UE8M0 scales, M = 128, N = 8, scale byte kb. +__host__ __device__ constexpr uint32_t qkv_idesc(int kb) +{ + return (static_cast(kb) << 4) | ((8u >> 3) << 17) | (1u << 23) | ((128u >> 7) << 27) + | (static_cast(kb) << 29); +} + +__device__ __forceinline__ void umma_mx( + uint32_t tmem_d, uint32_t tmem_a, uint64_t bdesc, uint32_t idesc, uint32_t sfa, uint32_t sfb, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" + " tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [%0], [%1], %2, %3, [%4], [%5], p;\n}\n" :: + "r"(tmem_d), + "r"(tmem_a), "l"(bdesc), "r"(idesc), "r"(sfa), "r"(sfb), "r"(accumulate)); +} + +__device__ __forceinline__ float pow2_from_ue8m0(uint32_t e) +{ + return __uint_as_float(e == 0 ? 0x00400000u : (e << 23)); +} + +// Unscaled E4M3 x E4M3 MMA, FP32 accumulator: M = 128, N = 8, K = 32, a fresh accumulator. +constexpr uint32_t kQkvIdescF8 = (1u << 4) | ((8u >> 3) << 17) | ((128u >> 4) << 24); + +__device__ __forceinline__ void umma_f8_ss(uint32_t tmem_d, uint64_t adesc, uint64_t bdesc) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, 0, 0;\n tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(kQkvIdescF8)); +} + +__device__ __forceinline__ void tmem_ld16_nowait(uint32_t addr, float* v) +{ + uint32_t r[16]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, " + "%15}, [%16];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]) + : "r"(addr)); +#pragma unroll + for (int i = 0; i < 16; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +__device__ __forceinline__ void umma_mx_ss( + uint32_t tmem_d, uint64_t adesc, uint64_t bdesc, uint32_t idesc, uint32_t sfa, uint32_t sfb, int accumulate) +{ + asm volatile( + "{\n .reg .pred p;\n setp.ne.b32 p, %6, 0;\n" + " tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [%0], %1, %2, %3, [%4], [%5], p;\n}\n" :: + "r"(tmem_d), + "l"(adesc), "l"(bdesc), "r"(idesc), "r"(sfa), "r"(sfb), "r"(accumulate)); +} + +__device__ __forceinline__ bool elect_one() +{ + uint32_t pred; + asm volatile("{\n .reg .pred P;\n elect.sync _|P, 0xffffffff;\n selp.u32 %0, 1, 0, P;\n}\n" : "=r"(pred)); + return pred != 0; +} + +__device__ __forceinline__ int qkv_parity(Params const& p) +{ + int v; + asm volatile("ld.relaxed.gpu.global.s32 %0, [%1];\n" : "=r"(v) : "l"(p.counters) : "memory"); + return v & 1; +} + +// Barriers and TMEM (warp 1), then the first ring stages (warp 0). Warps 0 and 1 of every CTA call it. +template +__device__ __forceinline__ void qkv_setup( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int lane, int u) +{ + constexpr bool kPipe = M3FA_QKV_PIPE != 0; + if (kPipe && warp == 0 && lane == 0 && u >= 0) + { + // First thing at entry: this unit's whole weight slice. The issuing thread initializes the ring's full + // barriers itself; every waiter on them passes bar.sync 8 below, which this thread reaches after the init. + asm volatile("prefetch.tensormap [%0];\n" ::"l"(&p.wmap) : "memory"); + int const rt = u / kQkvSplits; + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + for (int i = 0; i < kQkvRing; ++i) + { + mbar_init(gbar + 8 * i, 1); + } + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + for (int i = 0; i < kQkvRing && i < nt; ++i) + { + uint32_t const full = gbar + 8 * i; + mbar_expect_tx(full, kQkvWTile + kQkvSfAtom); + if (p.wpre != nullptr) + { + // Pre-tiled weight: the tile's 16 KB shared-memory image is contiguous (one bulk copy). + bulk_g2s(smem + SX::kRingW + i * kQkvWTile, + p.wpre + (static_cast(rt) * kQkvKTiles + kt0 + i) * kQkvWTile, kQkvWTile, full); + } + else + { + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap, (kt0 + i) * 128, rt * 128, full); + } + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf + (static_cast(rt) * kQkvKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, full); + } + M3FA_PROF_T(34, 0); + } + if (warp == 1) + { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n" ::"r"(smem_u32(tmem_slot)), + "r"(kTmemAll) + : "memory"); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n" ::: "memory"); + M3FA_PROF_T(35, 32); + if (lane == 0) + { + for (int i = (kPipe && u >= 0) ? kQkvRing : 0; i < 2 * kQkvRing + 2 + kQkvMaxTiles; ++i) + { + // PIPE 2: tile_done[0] (unused by the scaled MMAs) is act_ready, one arrival per thread outside warp 1 + mbar_init(gbar + 8 * i, M3FA_QKV_PIPE == 2 && i == 2 * kQkvRing + 2 ? kThreads - 32 : 1); + } + // The split-K reduce barrier: this CTA's own arrival carries the peers' partial bytes. + uint32_t const red = gbar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); + mbar_init(red, 1); + mbar_expect_tx(red, (kQkvSplits - 1) * kQkvOwnRows * kMaxTok * 4); + uint32_t const ob = smem_u32(smem + SX::kOBar); + for (int i = 0; i < kOpTiles + 3; ++i) + { + mbar_init(ob + 8 * i, 1); + } + mbar_expect_tx(ob + 8 * (kOpTiles + 2), (kOpSplits - 1) * kOpOwnRows * kMaxTok * 4); + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + asm volatile("bar.sync 8, 64;\n" ::: "memory"); + if (!kPipe && warp == 0 && lane == 0 && u >= 0) + { + int const rt = u / kQkvSplits; + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + for (int i = 0; i < kQkvRing && i < nt; ++i) + { + uint32_t const full = gbar + 8 * i; + mbar_expect_tx(full, kQkvWTile + kQkvSfAtom); + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap, (kt0 + i) * 128, rt * 128, full); + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf + (static_cast(rt) * kQkvKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, full); + } + M3FA_PROF_T(34, 0); + } +} + +// The rest of the ring (warp 0, lane 0) and the copies into TMEM (warp 1, lane 0); before the PDL wait. +template +__device__ __forceinline__ void qkv_stream_lane( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int u) +{ + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + uint32_t const full = gbar; + uint32_t const a_ready = gbar + 16 * kQkvRing; + if (warp == 0 || !kQkvScaled) + { + return; // every tile was issued in qkv_setup; the unscaled MMAs read them in place + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + // Fully unrolled with compile-time offsets: tcgen05 operands stay in uniform registers (a per-instruction + // conversion of computed operands costs ~150 cycles per MMA or copy). + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint32_t const sf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + mbar_wait_cta(full + 8 * i, 0); + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if constexpr (!kQkvASmem) + { +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kTA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + } + tmem_cp_sf(tm + kTSfa + i * 4, sf0 + i * kQkvSfAtom); + } + } + M3FA_PROF_T(31, 32); + umma_commit(a_ready); +} + +__device__ __forceinline__ bool mbar_test(uint32_t bar, uint32_t parity) +{ + uint32_t ready; + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.shared::cta.b64 p, [%1], %2;\n selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(bar), "r"(parity) + : "memory"); + return ready != 0; +} + +// PIPE 2, warp 1's elected thread from before the PDL wait to the last MMA. Tiles that land before the activation +// is ready go on into TMEM (A and scale atom), off the critical path; once the activation is in (act_ready), the +// MMAs issue in tile order: the copied tiles with A in TMEM, every later tile as it lands with A read in place from +// the ring. The MMA sequence (and so the accumulation order) is the same either way. +template +__device__ __forceinline__ void qkv_hybrid(Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int u) +{ + int const kt0 = qkv_tile0(u % kQkvSplits); + int const nt = qkv_tile0(u % kQkvSplits + 1) - kt0; + uint32_t const act_ready = gbar + 8 * (2 * kQkvRing + 2); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + SX::kXA)); + uint32_t const wsf0 = smem_u32(smem + SX::kRingSf); + uint32_t const xsf0 = smem_u32(smem + SX::kXASf); + int copied = 0; // tiles [0, copied) are in TMEM + bool act = false; +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt && !act) + { + while (!mbar_test(gbar + 8 * i, 0)) + { + if (mbar_test(act_ready, 0)) + { + act = true; + break; + } + } + if (!act) + { + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kTA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + copied = i + 1; + } + } + } + while (!act) + { + act = mbar_test(act_ready, 0); + } + M3FA_PROF_T(21, 32); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + bool const in_tmem = i < copied; + if (!in_tmem) + { + mbar_wait_epi(gbar + 8 * i, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + } + uint32_t const sfb = tm + kTSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + uint64_t const bd = b0 + ((i * 1024 + kb * 32) >> 4); + if (in_tmem) + { + umma_mx(tm + kTAcc, tm + kTA + j * 8, bd, qkv_idesc(kb), tm + kTSfa + i * 4, sfb, j > 0); + } + else + { + umma_mx_ss(tm + kTAcc, w0 + ((i * kQkvWTile + kb * 32) >> 4), bd, qkv_idesc(kb), tm + kTSfa + i * 4, + sfb, j > 0); + } + } + } + } + M3FA_PROF_T(22, 32); + M3FA_PROF_T(31, 32); + umma_commit(gbar + 16 * kQkvRing); // a_ready (for its waiters) + umma_commit(gbar + 16 * kQkvRing + 8); // acc_done +} + +// The other lanes of warps 0 and 1 hold at __syncwarp rather than at griddepcontrol.wait, which would stall the +// whole warp (lane 0's loop included) until the producer grid completes. +template +__device__ __forceinline__ void qkv_stream( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int warp, int lane, int u) +{ + if (M3FA_QKV_PIPE == 2 && u >= 0 && warp == 1) + { + if (elect_one()) + { + qkv_hybrid(p, smem, gbar, tmem_slot, u); + } + __syncwarp(); + return; + } + if (u < 0 || warp > 1 || M3FA_QKV_PIPE) + { + return; // pipelined: the MMA thread takes each tile as it lands + } + if (lane == 0) + { + qkv_stream_lane(p, smem, gbar, tmem_slot, warp, u); + } + __syncwarp(); +} + +// After the PDL wait. Every thread quantizes a piece of the x slice into the B tiles (named barrier 9); warp 1 then +// issues the scale copies and the MMAs. +template +__device__ __forceinline__ void qkv_mma( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int tid, int u) +{ + int const warp = tid >> 5; + int const s = u % kQkvSplits; + int const kt0 = qkv_tile0(s); + int const nt = qkv_tile0(s + 1) - kt0; + { + // The MXFP8 activation slice (E4M3 rows of the unit's k-tiles in 128-byte-swizzled B tiles, and their scale + // atoms), copied with plain 16-byte loads over the overlaid ring once its tiles went on into TMEM: every load + // first, then the stores (one L2 round trip). + constexpr int kXDst = M3FA_QKV_PIPE ? SX::kXA : SX::kXt; + constexpr int kXSfDst = M3FA_QKV_PIPE ? SX::kXASf : SX::kXsf; + if (!M3FA_QKV_PIPE) + { + mbar_wait_cta(gbar + 16 * kQkvRing, 0); // a_ready: the ring's weight tiles are in TMEM + } + constexpr int kXItems = kQkvMaxTiles * 8 * kMaxTok; // (tile, 16-byte chunk, token) + constexpr int kSfItems = kQkvMaxTiles * (kQkvSfAtom / 16); + constexpr bool kHyb = M3FA_QKV_PIPE == 2; + constexpr int kCopiers = kHyb ? kThreads - 32 : kThreads; // PIPE 2: warp 1 is issuing the tiles + if (kHyb && warp == 1) + { + return; + } + int const ctid = kHyb && tid >= 64 ? tid - 32 : tid; + constexpr int kPer = (kXItems + kSfItems + kCopiers - 1) / kCopiers; + uint4 v[kPer]; + int dst[kPer]; +#pragma unroll + for (int k = 0; k < kPer; ++k) + { + int const q = ctid + kCopiers * k; + dst[k] = -1; + if (q < kXItems) + { + int const t = q % kMaxTok; + int const c = (q / kMaxTok) % 8; + int const i = q / (kMaxTok * 8); + if (i < nt && t < p.num_tokens) + { + v[k] = *reinterpret_cast( + p.xq + static_cast(t) * p.ldxq + (kt0 + i) * 128 + c * 16); + dst[k] = kXDst + i * 1024 + t * 128 + ((c ^ t) * 16); + } + } + else if (q < kXItems + kSfItems) + { + int const j = q - kXItems; + int const i = j / (kQkvSfAtom / 16); + int const c = j % (kQkvSfAtom / 16); + if (i < nt) + { + v[k] = *reinterpret_cast( + p.xsf_act + static_cast(kt0 + i) * kQkvSfAtom + c * 16); + dst[k] = kXSfDst + i * kQkvSfAtom + c * 16; + } + } + } +#pragma unroll + for (int k = 0; k < kPer; ++k) + { + if (dst[k] >= 0) + { + *reinterpret_cast(smem + dst[k]) = v[k]; + } + } + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + if (kHyb) + { + asm volatile("mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];\n" ::"r"(gbar + 8 * (2 * kQkvRing + 2)) + : "memory"); + return; + } + if (warp != 1) + { + asm volatile("bar.arrive 9, %0;\n" ::"n"(kThreads) : "memory"); + } + } + if (warp == 1) + { + M3FA_PROF_T(32, 32); + asm volatile("bar.sync 9, %0;\n" ::"n"(kThreads) : "memory"); + M3FA_PROF_T(21, 32); + // Operands computed by the whole warp (warp-uniform, kept in uniform registers); one elected lane issues. + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm = *tmem_slot; + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + (M3FA_QKV_PIPE ? SX::kXA : SX::kXt))); + uint64_t const a0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + if (elect_one()) + { + M3FA_CLK(41); + if constexpr (!kQkvScaled) + { +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + mbar_wait_cta(gbar + 8 * i, 0); // weight tile i landed + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + umma_f8_ss(tm + kTAcc + (i * 4 + kb) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4), + b0 + ((i * 1024 + kb * 32) >> 4)); + } + umma_commit(gbar + 8 * (2 * kQkvRing + 2 + i)); // tile i's blocks are in TMEM + } + } + } + else + { + uint32_t const xsf0 = smem_u32(smem + (M3FA_QKV_PIPE ? SX::kXASf : SX::kXsf)); + uint32_t const wsf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + if constexpr (M3FA_QKV_PIPE != 0) + { + // Tile i as it lands: its A copy and scale atom into TMEM, then its MMAs (issue order). + mbar_wait_cta(gbar + 8 * i, 0); + if (i == 0) + { + M3FA_PROF_T(30, 32); + } + if (i == nt - 1) + { + M3FA_CLK(47); // the unit's last weight tile landed (clock64 builds) + } + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if constexpr (!kQkvASmem) + { +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b( + tm + kTA + (i * 4 + kb) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + } + tmem_cp_sf(tm + kTSfa + i * 4, wsf0 + i * kQkvSfAtom); + } + uint32_t const sfb = tm + kTSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + uint64_t const bd = b0 + ((i * 1024 + kb * 32) >> 4); + if constexpr (kQkvASmem) + { + umma_mx_ss(tm + kTAcc + (j % kQkvNacc) * 8, a0 + ((i * kQkvWTile + kb * 32) >> 4), bd, + qkv_idesc(kb), tm + kTSfa + i * 4, sfb, j >= kQkvNacc); + } + else + { + umma_mx(tm + kTAcc + (j % kQkvNacc) * 8, tm + kTA + j * 8, bd, qkv_idesc(kb), + tm + kTSfa + i * 4, sfb, j >= kQkvNacc); + } + } + } + } + } + M3FA_CLK(42); + if (M3FA_QKV_PIPE) + { + M3FA_PROF_T(31, 32); + umma_commit(gbar + 16 * kQkvRing); // a_ready (kept for its waiters: every tile is consumed) + } + umma_commit(gbar + 16 * kQkvRing + 8); // acc_done + } + M3FA_PROF_T(22, 32); + } +} + +// Warps 16..19 (TMEM lanes 32 (warp % 4) + lane = the unit's rows): partial rows out, the other parity reset. +template +__device__ __forceinline__ void qkv_epilogue( + Params const& p, uint8_t* smem, uint32_t gbar, uint32_t* tmem_slot, int tid, int u, int par, unsigned int calls) +{ + int const warp = tid >> 5; + int const lane = tid & 31; + if constexpr (kQkvScaled) + { + mbar_wait_epi(gbar + 16 * kQkvRing + 8, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + M3FA_PROF_T(24, 512); + if (tid == 512) + { + M3FA_CLK(43); + } + int const r = (warp & 3) * 32 + lane; // row within the unit's tile + uint32_t const taddr = *tmem_slot + (((warp & 3) * 32) << 16) + kTAcc; + float val[kMaxTok]; + if constexpr (kQkvScaled) + { + float v[32]; + tmem_ld16(taddr, *reinterpret_cast(v)); + if constexpr (kQkvNacc == 4) + { + tmem_ld16(taddr + 16, *reinterpret_cast(v + 16)); + } +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = kQkvNacc == 4 ? (v[t] + v[8 + t]) + (v[16 + t] + v[24 + t]) + : kQkvNacc == 2 ? v[t] + v[8 + t] + : v[t]; + } + } + else + { + // Sum over K blocks in order: block result x (weight scale x activation scale). + int const nt = qkv_tile0(u % kQkvSplits + 1) - qkv_tile0(u % kQkvSplits); + float const* xsc = reinterpret_cast(smem + SX::kXsc); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = 0.0f; + } +#pragma unroll + for (int i = 0; i < kQkvMaxTiles; ++i) + { + if (i < nt) + { + // Tile i's four blocks as soon as its MMAs complete: the epilogue trails the MMAs by one tile. + mbar_wait_cta(gbar + 8 * (2 * kQkvRing + 2 + i), 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + float d[32]; + tmem_ld16_nowait(taddr + i * 32, d); + tmem_ld16_nowait(taddr + i * 32 + 16, d + 16); + uint32_t const wsw = *reinterpret_cast( + smem + SX::kRingSf + i * kQkvSfAtom + (r % 32) * 16 + (r / 32) * 4); + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + float const ws = pow2_from_ue8m0((wsw >> (8 * kb)) & 0xFFu); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + val[t] = fmaf(d[kb * 8 + t], ws * xsc[(i * 4 + kb) * kMaxTok + t], val[t]); + } + } + } + } + } + // v52's split-K reduce: split s owns rows [32 s, 32 s + 32) of the tile; the other splits push their fp32 partials + // of those rows into its mailbox (st.async, completing its reduce barrier's bytes); the owner adds the splits in + // rank order from 0.0 and stores bf16. + int const s = u % kQkvSplits; + int const owner = r / kQkvOwnRows; + int const local = r % kQkvOwnRows; + uint32_t const red = gbar + 8 * (2 * kQkvRing + 2 + kQkvMaxTiles); + uint32_t cta_rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + // Units are numbered down from the last CTA: split q of this tile is CTA rank cta_rank - (q - s). + if (owner != s) + { + uint32_t const peer = cta_rank + s - owner; + uint32_t const dst + = mapa_shared(smem_u32(smem + SX::kMailbox) + ((s * kQkvOwnRows + local) * kQkvMbStride) * 4, peer); + uint32_t const pbar = mapa_shared(red, peer); +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];\n" ::"r"( + dst + 16 * g), + "r"(__float_as_uint(val[4 * g])), "r"(__float_as_uint(val[4 * g + 1])), + "r"(__float_as_uint(val[4 * g + 2])), "r"(__float_as_uint(val[4 * g + 3])), "r"(pbar) + : "memory"); + } + } + else + { + M3FA_PROF_T(26, tid - lane); + // test_wait (no suspend): completions by peers' st.async bytes wake a suspended try_wait late. + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(red) + : "memory"); + } + M3FA_PROF_T(27, tid - lane); + float const* mb = reinterpret_cast(smem + SX::kMailbox); + int const row = (u / kQkvSplits) * 128 + r; + // Lamport hand-off: this call's rows go to buffer (calls & 1) with relaxed stores, the consumers poll the + // values themselves (0xFFFF, a NaN the rounding of a finite sum never produces, marks "not yet"); the other + // buffer, last read by the previous call (complete: this grid passed its PDL wait), is reset to the marker. + uint16_t* cur = reinterpret_cast(p.qkv_out) + (calls & 1u) * kMaxTok * kQkvN; + uint16_t* nxt = reinterpret_cast(p.qkv_out) + ((calls & 1u) ^ 1u) * kMaxTok * kQkvN; +#pragma unroll + float qsum[kMaxTok]; +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + float4 part[kQkvSplits]; +#pragma unroll + for (int q = 0; q < kQkvSplits; ++q) + { + part[q] = *reinterpret_cast(mb + (q * kQkvOwnRows + local) * kQkvMbStride + 4 * g); + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float sum = 0.0F; +#pragma unroll + for (int q = 0; q < kQkvSplits; ++q) + { + float const pj = j == 0 ? part[q].x : j == 1 ? part[q].y : j == 2 ? part[q].z : part[q].w; + sum = sum + (q == s ? val[4 * g + j] : pj); + } + qsum[4 * g + j] = sum; + } + } +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + float const sum = qsum[t]; + if (t < p.num_tokens) + { + __nv_bfloat16 const b = __float2bfloat16_rn(sum); + uint16_t bits = *reinterpret_cast(&b); + // Never the "not yet" marker: a 0xFFFF NaN sum is stored as the NaN 0x7FFF. Only a NaN projection row + // changes (its K/V cache bytes then differ from v52 + p6 in the NaN's encoding); finite rows are exact. + bits = bits == 0xFFFFu ? static_cast(0x7FFFu) : bits; +#ifndef M3FA_QKV_REL +#define M3FA_QKV_REL 0 // 1: the row store is a release (tried: whether it becomes visible to the pollers sooner) +#endif + if (M3FA_QKV_REL) + { + asm volatile( + "st.release.gpu.global.u16 [%0], %1;\n" ::"l"(cur + static_cast(t) * kQkvN + row), + "h"(bits) + : "memory"); + } + else + { + asm volatile( + "st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(cur + static_cast(t) * kQkvN + row), + "h"(bits) + : "memory"); + } + } + asm volatile("st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(nxt + static_cast(t) * kQkvN + row), + "h"(static_cast(0xFFFFu)) + : "memory"); + } + M3FA_PROF_T(28, tid - lane); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + M3FA_PROF_T(25, 512); + if (tid == 512) + { + M3FA_CLK(46); + } +} + +// ---------------- Megakernel o_proj (v52 m3_mx_gemv bit for bit: split 2, contiguous k-tiles, nacc 1) ---------------- +// The unit's weight tiles and scale atoms into the ring (warp 0, lane 0). From L2 when the CTA prefetched them there. +template +__device__ __forceinline__ void op_issue(Params const& p, uint8_t* smem, int ou) +{ + int const rt = ou / kOpSplits; + int const kt0 = (ou % kOpSplits) * kOpTiles; + uint32_t const ob = smem_u32(smem + SX::kOBar); + for (int i = 0; i < kOpTiles; ++i) + { + mbar_expect_tx(ob + 8 * i, kQkvWTile + kQkvSfAtom); + tma_load_2d(smem_u32(smem + SX::kRingW + i * kQkvWTile), &p.wmap_o, (kt0 + i) * 128, rt * 128, ob + 8 * i); + bulk_g2s(smem + SX::kRingSf + i * kQkvSfAtom, + p.wsf_o + (static_cast(rt) * kOpKTiles + kt0 + i) * kQkvSfAtom, kQkvSfAtom, ob + 8 * i); + } +} + +// L2 prefetch of the unit's weight tiles and scale atoms (one lane; nothing waits for it). +// The QKV units' share of the window-1 L2 prefetch (the next consumer's weights), issued once their own QKV weight +// stream has landed (acc_done): the attention core reads little from HBM. +__device__ __forceinline__ void window_prefetch_l2(Params const& p, int u) +{ + constexpr long long kChunk = 32768; + long long const n0 = (p.pf_b0 + kChunk - 1) / kChunk, n1 = (p.pf_b1 + kChunk - 1) / kChunk; + uint64_t pol = 0; + if (p.pf_hint) + { + asm volatile("createpolicy.fractional.L2::evict_last.b64 %0, 1.0;\n" : "=l"(pol)); + } + for (long long c = u; c < n0 + n1; c += kQkvUnits) + { + long long const a = c < n0 ? p.pf_a0 + c * kChunk : p.pf_a1 + (c - n0) * kChunk; + long long const rem = c < n0 ? p.pf_b0 - c * kChunk : p.pf_b1 - (c - n0) * kChunk; + unsigned int const n = static_cast(rem < kChunk ? rem : kChunk) & ~15u; + if (n == 0) + { + continue; + } + if (p.pf_hint) + { + asm volatile("cp.async.bulk.prefetch.L2.global.L2::cache_hint [%0], %1, %2;\n" ::"l"(a), "r"(n), "l"(pol) + : "memory"); + } + else + { + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;\n" ::"l"(a), "r"(n) : "memory"); + } + } +} + +__device__ __forceinline__ void op_prefetch_l2(Params const& p, int ou) +{ + int const rt = ou / kOpSplits; + int const kt0 = (ou % kOpSplits) * kOpTiles; + for (int i = 0; i < kOpTiles; ++i) + { + asm volatile("cp.async.bulk.prefetch.tensor.2d.L2.global [%0, {%1, %2}];\n" ::"l"(&p.wmap_o), + "r"((kt0 + i) * 128), "r"(rt * 128) + : "memory"); + } + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;\n" ::"l"( + p.wsf_o + (static_cast(rt) * kOpKTiles + kt0) * kQkvSfAtom), + "r"(kOpTiles * kQkvSfAtom) + : "memory"); +} + +// The weight tiles into TMEM as they land (warp 1, elected lane, operands warp-uniform). +template +__device__ __forceinline__ void op_stream(uint8_t* smem, uint32_t tm) +{ + uint32_t const ob = smem_u32(smem + SX::kOBar); + uint64_t const w0 = umma_desc_sw128(smem_u32(smem + SX::kRingW)); + uint32_t const sf0 = smem_u32(smem + SX::kRingSf); +#pragma unroll + for (int i = 0; i < kOpTiles; ++i) + { + mbar_wait_cta(ob + 8 * i, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + tmem_cp_128x256b(tm + kOA + (i * 4 + kb) * 8, w0 + ((i * kQkvWTile + kb * 32) >> 4)); + } + tmem_cp_sf(tm + kOSfa + i * 4, sf0 + i * kQkvSfAtom); + } + umma_commit(ob + 8 * kOpTiles); // a_ready +} + +// The attention's MXFP8 rows of this unit's k-tiles (the heads [kt0, kt0 + 8)), Lamport-polled value by value: 0xFF +// is NaN in both E4M3 and UE8M0, which the merge never writes (see the merge). Every thread of the CTA, one item each. +template +__device__ __forceinline__ void op_act(Params const& p, uint8_t* smem, int ou, int tid, unsigned int calls) +{ + int const kt0 = (ou % kOpSplits) * kOpTiles; + uint8_t const* q8 = p.q8 + (calls & 1u) * kMaxTok * kOpK; + uint8_t const* sf = p.sf + (calls & 1u) * kOpSfBytes; + constexpr int kXItems = kOpTiles * 8 * kMaxTok; + constexpr int kSfItems = kOpTiles * (kQkvSfAtom / 16); + for (int q = tid; q < kXItems + kSfItems; q += kThreads) + { + if (q < kXItems) + { + int const t = q % kMaxTok; + int const c = (q / kMaxTok) % 8; + int const i = q / (kMaxTok * 8); + if (t < p.num_tokens) + { + uint4 v; + bool pending; + do + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(q8 + static_cast(t) * kOpK + (kt0 + i) * 128 + c * 16) + : "memory"); + uint32_t const w[4] = {v.x, v.y, v.z, v.w}; + pending = false; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + // a byte equal to 0xFF: (x & 0x7F7F7F7F) + 0x01010101 sets bit 7 of exactly the 0x7F/0xFF bytes + uint32_t const hit = (((w[j] & 0x7F7F7F7Fu) + 0x01010101u) & w[j] & 0x80808080u); + pending |= hit != 0u; + } + } while (pending); + *reinterpret_cast(smem + SX::kOXt + i * 1024 + t * 128 + ((c ^ t) * 16)) = v; + } + } + else if (q < kXItems + kSfItems) + { + int const j = q - kXItems; + int const i = j / (kQkvSfAtom / 16); + int const c = j % (kQkvSfAtom / 16); + uint8_t const* src = sf + (kt0 + i) * kQkvSfAtom + c * 16; + uint4 v; + if (c < p.num_tokens) + { + // Row c of the atom: its 4 block scales in bytes [16 c, 16 c + 4). + do + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(src) + : "memory"); + } while ((((v.x & 0x7F7F7F7Fu) + 0x01010101u) & v.x & 0x80808080u) != 0u); + } + else + { + v = *reinterpret_cast(src); + } + *reinterpret_cast(smem + SX::kOXsf + i * kQkvSfAtom + c * 16) = v; + } + } +} + +// The MMAs (warp 1, elected lane): v52's order (per k-tile: B scale copy, 4 MMAs; one accumulator). +template +__device__ __forceinline__ void op_mma(uint8_t* smem, uint32_t tm) +{ + uint32_t const ob = smem_u32(smem + SX::kOBar); + mbar_wait_cta(ob + 8 * kOpTiles, 0); // a_ready + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint64_t const b0 = umma_desc_sw128(smem_u32(smem + SX::kOXt)); + uint32_t const xsf0 = smem_u32(smem + SX::kOXsf); +#pragma unroll + for (int i = 0; i < kOpTiles; ++i) + { + uint32_t const sfb = tm + kOSfb + i * 4; + tmem_cp_sf(sfb, xsf0 + i * kQkvSfAtom); +#pragma unroll + for (int kb = 0; kb < 4; ++kb) + { + int const j = i * 4 + kb; + umma_mx(tm + kOAcc, tm + kOA + j * 8, b0 + ((i * 1024 + kb * 32) >> 4), qkv_idesc(kb), tm + kOSfa + i * 4, + sfb, j > 0); + } + } + umma_commit(ob + 8 * (kOpTiles + 1)); // acc_done +} + +// Warps 16..19: the pair's split-K reduce (owner rows [64 s, 64 s + 64)) and the bf16 rows of y. +template +__device__ __forceinline__ void op_epilogue(Params const& p, uint8_t* smem, uint32_t tm, int tid, int ou) +{ + int const warp = tid >> 5; + int const lane = tid & 31; + uint32_t const ob = smem_u32(smem + SX::kOBar); + // The store's kernel parameters, fetched before the wait: a first touch of their constant-bank line costs a miss. + __nv_bfloat16* y_out = p.y_out; + int num_tokens = p.num_tokens; + int ar_mode = p.ar_mode; + asm volatile("" : "+l"(y_out), "+r"(num_tokens), "+r"(ar_mode)); +#ifndef M3FA_OP_PARPRE +#define M3FA_OP_PARPRE 1 // read the all-reduce slot parities before the accumulator wait (a global load off the tail) +#endif + // The parity flags are stable here: the all-reduce that flips them last completed before this grid's PDL wait, + // and the next one reads them only after this kernel's pushes. + int par_pre[kMaxTok]; + { + int const row_pre = (ou / kOpSplits) * 128 + ((tid >> 5) & 3) * 32 + (tid & 31); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + par_pre[t] = 0; + if (M3FA_OP_PARPRE && ar_mode != 0 && t < num_tokens) + { + asm volatile("ld.volatile.global.s32 %0, [%1];\n" + : "=r"(par_pre[t]) + : "l"(p.ar_flags + t * 8 + row_pre / (kOpN / 8)) + : "memory"); + } + } + } + mbar_wait_epi(ob + 8 * (kOpTiles + 1), 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF_T(51, kComputeThreads); + int const r = (warp & 3) * 32 + lane; + float v[16]; + tmem_ld16(tm + (((warp & 3) * 32) << 16) + kOAcc, v); + int const s = ou % kOpSplits; + int const owner = r / kOpOwnRows; + int const local = r % kOpOwnRows; + uint32_t const red = ob + 8 * (kOpTiles + 2); + uint32_t cta_rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + if (owner != s) + { + uint32_t const peer = cta_rank + owner - s; // o_proj units count up with the CTA index + uint32_t const dst + = mapa_shared(smem_u32(smem + SX::kOMailbox) + ((s * kOpOwnRows + local) * kQkvMbStride) * 4, peer); + uint32_t const pbar = mapa_shared(red, peer); +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];\n" ::"r"( + dst + 16 * g), + "r"(__float_as_uint(v[4 * g])), "r"(__float_as_uint(v[4 * g + 1])), "r"(__float_as_uint(v[4 * g + 2])), + "r"(__float_as_uint(v[4 * g + 3])), "r"(pbar) + : "memory"); + } + } + else + { + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(ready) + : "r"(red) + : "memory"); + } + M3FA_PROF_T(52, tid - lane); + float const* mb = reinterpret_cast(smem + SX::kOMailbox); + int const row = (ou / kOpSplits) * 128 + r; + float sums[kMaxTok]; +#pragma unroll + for (int g = 0; g < kMaxTok / 4; ++g) + { + float4 part[kOpSplits]; +#pragma unroll + for (int q = 0; q < kOpSplits; ++q) + { + part[q] = *reinterpret_cast(mb + (q * kOpOwnRows + local) * kQkvMbStride + 4 * g); + } +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float sum = 0.0F; +#pragma unroll + for (int q = 0; q < kOpSplits; ++q) + { + float const pj = j == 0 ? part[q].x : j == 1 ? part[q].y : j == 2 ? part[q].z : part[q].w; + sum = sum + (q == s ? v[4 * g + j] : pj); + } + sums[4 * g + j] = sum; + } + } + M3FA_PROF_T(54, tid - lane); +#pragma unroll + for (int t = 0; t < kMaxTok; ++t) + { + float const sum = sums[t]; +#ifndef M3FA_EXP +#define M3FA_EXP 0 +#endif + if (t < num_tokens) + { + if (M3FA_EXP == 1) + { + asm volatile( + "st.relaxed.gpu.global.u16 [%0], %1;\n" ::"l"(y_out + static_cast(t) * kOpN + row), + "h"(__bfloat16_as_ushort(__float2bfloat16_rn(sum))) + : "memory"); + } + else + { + y_out[static_cast(t) * kOpN + row] = __float2bfloat16_rn(sum); + } + } + if (ar_mode != 0) + { + // m3_mx_gemv's push into the all-reduce exchange: lane pairs hold row pairs; the even lane writes the + // bf16x2 word into this rank's slot of the (token, column slice) parity the AR's flags select. + float const hi = __shfl_xor_sync(0xffffffffu, sum, 1); + if ((lane & 1) == 0 && t < num_tokens) + { + int par; + if (M3FA_OP_PARPRE) + { + par = par_pre[t]; + } + else + { + asm volatile("ld.volatile.global.s32 %0, [%1];\n" + : "=r"(par) + : "l"(p.ar_flags + t * 8 + row / (kOpN / 8)) + : "memory"); + } + int64_t const off = static_cast(par & 1) * p.ar_maxt * p.ar_world * (kOpN / 2) + + (static_cast(t) * p.ar_world + p.ar_rank) * (kOpN / 2) + row / 2; + __nv_bfloat162 w2 = __floats2bfloat162_rn(sum, hi); + uint32_t word = *reinterpret_cast(&w2); + // -0.0 halves to +0.0: 0x80000000 marks an empty exchange word (as m3_mx_gemv's _bf16x2_word) + word = (word & 0xFFFFu) == 0x8000u ? (word & 0xFFFF0000u) : word; + word = (word >> 16) == 0x8000u ? (word & 0x0000FFFFu) : word; + asm volatile("st.global.b32 [%0], %1;\n" ::"l"(p.ar_ws + off), "r"(word) : "memory"); + } + } + } + M3FA_PROF_T(53, tid - lane); + } + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); +} + +// Packed-QKV element group [token][col * 128 + 4 * lane, +4): the split partials, polled until every one is stored, +// summed in split order and rounded to BF16. +__device__ __forceinline__ void qkv_load(Params const& p, int par, int tok, int col, int lane, float (&e)[4]) +{ + float const* base = p.part + par * kQkvPartStride + static_cast(tok) * kQkvN + col * kHeadDim + lane * 4; + uint4 pr[kQkvSplits]; + long long spins = 0; + bool ready; + do + { + ready = true; +#pragma unroll + for (int s = 0; s < kQkvSplits; ++s) + { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(pr[s].x), "=r"(pr[s].y), "=r"(pr[s].z), "=r"(pr[s].w) + : "l"(base + static_cast(s) * kMaxTok * kQkvN) + : "memory"); + ready &= pr[s].x != kQkvSentinel && pr[s].y != kQkvSentinel && pr[s].z != kQkvSentinel + && pr[s].w != kQkvSentinel; + } + if (++spins > (1LL << 26)) + { + __trap(); + } + } while (!ready); + float s0 = __uint_as_float(pr[0].x), s1 = __uint_as_float(pr[0].y), s2 = __uint_as_float(pr[0].z), + s3 = __uint_as_float(pr[0].w); +#pragma unroll + for (int s = 1; s < kQkvSplits; ++s) + { + s0 += __uint_as_float(pr[s].x); + s1 += __uint_as_float(pr[s].y); + s2 += __uint_as_float(pr[s].z); + s3 += __uint_as_float(pr[s].w); + } + e[0] = __bfloat162float(__float2bfloat16_rn(s0)); + e[1] = __bfloat162float(__float2bfloat16_rn(s1)); + e[2] = __bfloat162float(__float2bfloat16_rn(s2)); + e[3] = __bfloat162float(__float2bfloat16_rn(s3)); +} + +// Every CTA once it has read the call parity and the partials it needs: the last one flips the parity. +__device__ __forceinline__ void qkv_done(Params const& p) +{ + if (atomicAdd(p.counters + 1, 1) == kQkvGrid - 1) + { + atomicExch(p.counters + 1, 0); + atomicXor(p.counters, 1); + } +} + +} // namespace + +namespace +{ + +// FQKV: the kernel also computes the QKV(+index) projection from the layer input (see qkv_mma) and runs on a fixed +// kQkvGrid-CTA grid; the first attn_clusters clusters then run the attention on the projection's output. +template +__global__ void __launch_bounds__(kThreads, 1) m3_fused_decode_attn_kernel(__grid_constant__ Params const p) +{ +#if MINIMAX_M3_ATTN_DEVICE_ENABLED + using S = SmemX; + constexpr int kRows = S::kRows; + constexpr int kTok = S::kTok; + constexpr int kCs = kBlocks / BPC; // cluster size + constexpr int kGroups = QLEN / QG; // clusters per request + constexpr uint32_t kPushBytes = (kHeadDim + 1) * 4; // one partial row: 128 outputs and the LSE + constexpr int kTasks = kHq * QG + 3 * QLEN; // producer tasks: 16 Q heads x QG, [K, V, index-K] x QLEN + constexpr int kMaxT = (kTasks + kWarps - 1) / kWarps; + constexpr int kChunks = BPC * kPage * 4; // 16-byte packed chunks per K (or V) tile set + constexpr int kKPer = (kChunks + kComputeThreads - 1) / kComputeThreads; // K chunks per thread [0, 512) + constexpr bool kRegCore = RegCore::kOn; // warp-MMA attention core (see RegCore) + // V tiles (raw copy before the PDL wait, dequantized during QK and the softmax): on the packed path every warp + // outside the softmax groups takes part (at M=1, warps 4..19), on the fused-projection path the producer warps. +#ifndef M3FA_RS2 +#define M3FA_RS2 1 // megakernel at one 16-row chunk: the softmax and push rows split over two warp groups (8 each) +#endif + // Split softmax: the megakernel and multi-row (spec verify) steps; at one-row decode on the packed path it + // measured slower (it takes four warps from the V dequantization there). + constexpr bool kRs2 = M3FA_RS2 && BPC * QG == 1 && !RegCore::kOn && (FQKV || QLEN > 1); +#ifndef M3FA_MEGA_WARPMMA +#define M3FA_MEGA_WARPMMA 1 +#endif + // Warp-issued MMAs and the early cluster barrier: measured gains in the megakernel and on multi-row / paired-block + // steps, but not at one-row decode on 16-CTA clusters (the packed path's (1, 1, 1)), where they stay off. + constexpr bool kTuned = FQKV || QLEN > 1 || BPC > 1; + constexpr bool kWarpMma = M3FA_MEGA_WARPMMA && kTuned; + constexpr bool kEarlyClw = M3FA_MEGA_EARLYCLW && kTuned; + constexpr int kSoftmaxGroups = kRs2 ? 2 : (BPC * QG) < 4 ? (BPC * QG) : 4; +#ifndef M3FA_MEGA_VWIDE +#define M3FA_MEGA_VWIDE \ + 1 // megakernel attention CTAs (no projection unit): every warp outside the softmax copies and + // dequantizes V, as on the packed path (not only the four producer warps) +#endif + constexpr int kVStart = (kRegCore || (FQKV && !M3FA_MEGA_VWIDE)) ? kComputeThreads : 128 * kSoftmaxGroups; + constexpr int kVThreads = kThreads - kVStart; // V chunks: the warps outside the softmax groups + constexpr int kVPer = (kChunks + kVThreads - 1) / kVThreads; + constexpr int kPairs = BPC * QG; // (block, 16-row chunk) pairs of the softmax and push + constexpr int kPairsPerGroup = (kPairs + 3) / 4; // per group of four compute warps + // Named barrier 6: the softmax warps (P written) and the producer warps (V written) before the PV issue. + constexpr int kPvBarThreads = 128 * kSoftmaxGroups + kVThreads; + extern __shared__ __align__(128) uint8_t smem[]; + float* row_scale = reinterpret_cast(smem + S::kRowScale); + float* row_lse = reinterpret_cast(smem + S::kRowLse); + float* red_m = reinterpret_cast(smem + S::kRedM); + float* red_l = reinterpret_cast(smem + S::kRedL); + float* recv = reinterpret_cast(smem + S::kRecv); + uint32_t const mbar = smem_u32(smem + S::kMbar); + uint32_t const mma_bar = mbar + 8; + uint32_t* tmem_slot = reinterpret_cast(smem + S::kMbar + 16); + + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const cr = blockIdx.x % kCs; // cluster rank; owns blocks [cr * BPC, cr * BPC + BPC) + int const ci = blockIdx.x / kCs; + int const req = ci / kGroups; + int const tfirst = (ci % kGroups) * QG; // this cluster's query tokens [tfirst, tfirst + QG) of the request + int const tok0 = req * QLEN; + uint32_t const gbar = smem_u32(smem + S::kGBar); // FQKV: weight ring, A ready, accumulators done + int const qkv_u = FQKV ? qkv_unit() : -1; // FQKV: this CTA's projection unit, or -1 + + // ---------------- Prologue: nothing here depends on the producer kernel ---------------- +#if M3FA_TRIGGER_AT == 0 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (!FQKV) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + M3FA_PROF(0); +#ifdef M3FA_PROFILE + if (p.prof != nullptr && tid == 0) + { + uint32_t smid; + asm volatile("mov.u32 %0, %%smid;" : "=r"(smid)); + p.prof[blockIdx.x * 64 + 40] = smid; + } +#endif + if constexpr (FQKV) + { + // The projection weights stream into TMEM while the producer (the all-reduce) still runs. + if (warp < 2) + { + qkv_setup(p, smem, gbar, tmem_slot, warp, lane, qkv_u); + } + if (ci >= p.attn_clusters) + { + // A projection-only CTA. Every reduce barrier of the cluster is initialized before any peer pushes. + asm volatile("barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;\n" :: + : "memory"); + int const op_u = static_cast(blockIdx.x) - kOpFirstCta; // this CTA's o_proj unit (>= 0 here) +#ifndef M3FA_OP_EARLY +#define M3FA_OP_EARLY 0 // 1: o_proj-only CTAs stream their weights from launch (0: after the PDL wait) +#endif + if (M3FA_OP_EARLY && qkv_u < 0 && warp == 0 && lane == 0) + { + op_issue(p, smem, op_u); // an o_proj-only CTA: its weights stream from launch on + } + if (M3FA_OP_L2PF && qkv_u >= 0 && warp == 0 && lane == 0) + { + op_prefetch_l2(p, op_u); // into L2 during the all-reduce, behind this CTA's QKV ring + } + qkv_stream(p, smem, gbar, tmem_slot, warp, lane, qkv_u); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (M3FA_MEGA_TRIGGER == 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + // This CTA's call counter: read after the PDL wait (the previous call's CTA increments it at its end), and + // early, so its latency stays off the epilogue's critical path. + unsigned int const calls = p.flags[32 + blockIdx.x]; + M3FA_PROF(1); + if (!M3FA_OP_EARLY && qkv_u < 0 && warp == 0 && lane == 0) + { + op_issue(p, smem, op_u); + } + if (qkv_u >= 0) + { + int const par = 0; + qkv_mma(p, smem, gbar, tmem_slot, tid, qkv_u); + if (warp >= kComputeWarps) + { + qkv_epilogue(p, smem, gbar, tmem_slot, tid, qkv_u, par, calls); + } + if (warp == 0 && lane == 0) + { + // The ring (and the activation over it) is free once the QKV MMAs completed. + mbar_wait_cta(gbar + 16 * kQkvRing + 8, 0); + op_issue(p, smem, op_u); + if (p.pf_b0 > 0 || p.pf_b1 > 0) + { + window_prefetch_l2(p, qkv_u); + } + } + } + __syncthreads(); // the QKV accumulator reads are done: TMEM columns [0, 512) are the o_proj's + M3FA_PROF(2); + uint32_t const tm = *tmem_slot; + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if (elect_one()) + { + op_stream(smem, tm); + } + __syncwarp(); + } + op_act(p, smem, op_u, tid, calls); + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + M3FA_PROF_T(36, 0); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Late release: the dependents launch once every unit has its attention input (the rel_v4 rule). + if (M3FA_MEGA_TRIGGER < 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + __syncthreads(); + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + if (elect_one()) + { + op_mma(smem, tm); + M3FA_PROF_T(50, threadIdx.x); + } + __syncwarp(); + } + if (warp >= kComputeWarps) + { + op_epilogue(p, smem, tm, tid, op_u); + } + M3FA_PROF_T(37, kComputeThreads); + __syncthreads(); + M3FA_PROF_T(55, 0); + if (M3FA_MEGA_TRIGGER == 3) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); // after this CTA's y rows / pushes + } + if (warp == 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" ::"r"(*tmem_slot), "r"(kTmemAll) + : "memory"); + M3FA_PROF_T(56, 32); + } + if (tid == 0) + { + p.flags[32 + blockIdx.x] = calls + 1u; // a plain store: no load on the exit path + } + M3FA_PROF(39); +#ifndef M3FA_FINAL_CLBAR +#define M3FA_FINAL_CLBAR 0 +#endif + if (M3FA_FINAL_CLBAR) + { + // No peer may exit while its partials can still be pushed into it (DSMEM lifetime). + asm volatile("barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;\n" :: + : "memory"); + } + M3FA_PROF(38); + return; + } + } +#ifndef M3FA_MEGA_ATTN_WAIT_FIRST +#define M3FA_MEGA_ATTN_WAIT_FIRST 0 // 1: the attention CTAs read nothing (metadata, caches) before their PDL wait +#endif +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (FQKV && M3FA_MEGA_ATTN_WAIT_FIRST) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } +#endif + if (!FQKV && warp == kWarps - 1) // a producer warp: it has nothing else to do before the PDL wait + { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n" ::"r"(smem_u32(tmem_slot)), + "r"(S::kTmemCols) + : "memory"); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n" ::: "memory"); + } + int const seq_len = p.seq_lens[req]; + int page[BPC]; +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + int const b = cr * BPC + h; + page[h] = b < p.bt_cols ? p.block_table[static_cast(req) * p.bt_stride + b] : 0; + } + // New-token cache slots, for the tokens up to this cluster's last one. + int slot[QLEN]; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + slot[t] = t < tfirst + QG ? p.out_cache_loc[tok0 + t] : -1; + } + // This warp's producer tasks and, for every lane that applies RoPE, the task token's cos/sin. + float cs[kMaxT][4], sn[kMaxT][4]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + int const t = i < kHq * QG ? tfirst + i / kHq : (i - kHq * QG) / 3; + if (i < kTasks && lane < 16 && t < QLEN) + { + int64_t const row = static_cast(p.position_ids[tok0 + t]) * kRotDim; +#pragma unroll + for (int e = 0; e < 4; ++e) + { + int const coefficient = (lane * 4 + e) & (kRotDim / 2 - 1); + cs[k][e] = p.rope[row + coefficient]; + sn[k][e] = p.rope[row + kRotDim / 2 + coefficient]; + } + } + else + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + cs[k][e] = 1.0F; + sn[k][e] = 0.0F; + } + } + } + float qw[4], kw[4], ikw[4]; + load_w4(qw, p.q_w, lane); + load_w4(kw, p.k_w, lane); + load_w4(ikw, p.ik_w, lane); + float const kqs = p.kv_quant_scale[1]; + float const vqs = p.kv_quant_scale[2]; + float const sl2e = (p.sm_scale * 1.4426950409f) * p.k_gs[0]; + float const vgs = p.v_gs[0]; + // Longest attended length among this cluster's tokens, and each block's rows within it. + int const kv_max = max(seq_len - QLEN + tfirst + QG, 0); + int nvh[BPC]; +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + nvh[h] = min(max(kv_max - (cr * BPC + h) * kPage, 0), kPage); + } + bool const any_valid = nvh[0] > 0; + if (tid == 32) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(mbar), "r"(1)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(mma_bar), "r"(1)); + asm volatile( + "mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(mbar), "r"(kBlocks * QG * BPC * kPushBytes) + : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + } + // This thread's packed K / V chunks, held in registers until they are dequantized after the PDL wait. + RawChunk kraw[kKPer]; +#pragma unroll + for (int k = 0; k < kKPer; ++k) + { + int const q = tid + kComputeThreads * k; + kraw[k] = RawChunk{make_uint4(0u, 0u, 0u, 0u), 0u}; + int const h = q / (kPage * 4); + if (q < kChunks && tid < kComputeThreads && nvh[h] > 0) + { + kraw[k] = load_chunk(p.kv_data + static_cast(page[h]) * p.d_page, + p.kv_scale + static_cast(page[h]) * p.s_page, chunk_row(q % (kPage * 4)), + chunk_col(q % (kPage * 4)), false); + } + } + // V goes through cp.async into shared memory: an asynchronous copy does not hold the producer warps up the way + // an outstanding register load does, and they need V only after QK has been issued. + if (tid >= kVStart) + { +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + if (nvh[h] > 0) + { + uint8_t* dst = smem + S::kRawV + h * S::kRawVBlock; + uint8_t const* data = p.kv_data + static_cast(page[h]) * p.d_page + p.d_plane; + uint8_t const* scales = p.kv_scale + static_cast(page[h]) * p.s_page + p.s_plane; + for (int c = tid - kVStart; c < S::kRawVBlock / 16; c += kVThreads) + { + cp_async16(dst + c * 16, c < kPage * 4 ? data + c * 16 : scales + (c - kPage * 4) * 16); + } + } + } + } + asm volatile("cp.async.commit_group;\n" ::: "memory"); + if (warp == kWarps - 1) + { + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + // Every CTA's mbarrier is initialized before any peer pushes into it; the matching wait sits right before the + // push, so this barrier's latency hides behind the attention (with aliased partials it moves there). + if constexpr (!S::kRecvAlias) + { + asm volatile("barrier.cluster.arrive.relaxed.aligned;\n" ::: "memory"); + } + int kv_len[QG]; +#pragma unroll + for (int tl = 0; tl < QG; ++tl) + { + kv_len[tl] = max(seq_len - QLEN + tfirst + tl + 1, 0); + } + int const cta_tok0 = cr * kTok; // request position of this CTA's first block row + // New token t's row in this CTA's tiles (block * 128 + row), or -1: the CTA owns the block holding the token's + // position, and the token is attended by this cluster (t <= its last token). Every thread derives the same. + int new_row[QLEN]; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + int const pos = seq_len - QLEN + t; + int const blk = min(pos >> 7, kBlocks - 1); + int const spage = slot[t] / kPage; + bool const mine = t < tfirst + QG && pos >= 0 && blk / BPC == cr && slot[t] >= 0 && spage < p.num_pages; + new_row[t] = mine && spage == page[blk % BPC] ? (blk % BPC) * kPage + (slot[t] & (kPage - 1)) : -1; + } + auto is_new_row = [&](int r) + { + bool hit = false; +#pragma unroll + for (int t = 0; t < QLEN; ++t) + { + hit |= new_row[t] == r; + } + return hit; + }; + // Cache writes of the new K/V/index-K are issued after this CTA has pushed its partials: only the fp16 tile rows + // are on the attention's path. kind 0: none, 1: NVFP4 data + scale, 2: FP8 index-K. + int wr_kind[kMaxT]; + uint8_t* wr_ptr[kMaxT]; + uint8_t* wr_sptr[kMaxT]; + uint32_t wr_val[kMaxT]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + wr_kind[k] = 0; + wr_ptr[k] = nullptr; + wr_sptr[k] = nullptr; + wr_val[k] = 0u; + } + uint32_t tmem = 0; + auto dequant_k = [&]() + { + if (tid < kComputeThreads) + { +#pragma unroll + for (int k = 0; k < kKPer; ++k) + { + int const q = tid + kComputeThreads * k; + if (q < kChunks) + { + int const h = q / (kPage * 4); + int const row = chunk_row(q % (kPage * 4)); + if (!is_new_row(h * kPage + row)) + { + store_chunk_fp16( + smem + S::kK16 + h * S::kTileBytes, row, chunk_col(q % (kPage * 4)), kraw[k], row < nvh[h]); + } + } + } + } + }; + // The producer warps' V tiles (raw tiles from cp.async, dequantized to fp16), then their PV-barrier arrival. + auto dequant_v = [&]() + { +#ifdef M3FA_DBG_NO_VDQ + if constexpr (!kRegCore) + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + return; +#endif + asm volatile("cp.async.wait_all;\n" ::: "memory"); + asm volatile("bar.sync 4, %0;\n" ::"n"(kVThreads) : "memory"); +#pragma unroll + for (int k = 0; k < kVPer; ++k) + { + int const q = tid - kVStart + kVThreads * k; + if (q < kChunks) + { + int const h = q / (kPage * 4); + int const row = chunk_row(q % (kPage * 4)); + int const cc = chunk_col(q % (kPage * 4)); + if (!is_new_row(h * kPage + row)) + { + uint8_t const* rawv = smem + S::kRawV + h * S::kRawVBlock; + uint8_t const* sc = rawv + kPage * 64 + (row >> 2) * 32 + (row & 3); + RawChunk c; + c.data = *reinterpret_cast(rawv + row * 64 + cc * 16); + c.scales + = static_cast(sc[(cc * 2) * 4]) | (static_cast(sc[(cc * 2 + 1) * 4]) << 8); + store_chunk_fp16(smem + S::kV16 + h * S::kTileBytes, row, cc, c, row < nvh[h]); + } + } + } + M3FA_FENCE_PROXY_ASYNC(); // V tiles -> tensor core reads (ordered by the barrier before the PV issue) + if constexpr (!kRegCore) + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + M3FA_PROF_T(16, kComputeThreads); + M3FA_PROF_T(17, kThreads - 1); + }; + if constexpr (FQKV) + { + qkv_stream(p, smem, gbar, tmem_slot, warp, lane, qkv_u); + } + if constexpr (!FQKV) + { + if (M3FA_V15_EARLYV && any_valid && tid >= kVStart) + { + dequant_v(); + } + if constexpr (kEarlyClw && !S::kRecvAlias) + { + asm volatile("barrier.cluster.wait.aligned;\n" ::: "memory"); + } + } + M3FA_PROF(11); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (FQKV && M3FA_MEGA_TRIGGER == 2) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#if M3FA_TRIGGER_AT == 1 + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +#endif +#endif + // FQKV (megakernel): this call's QKV rows sit in buffer (calls & 1), polled value by value below. + unsigned int const mega_calls = FQKV ? p.flags[32 + blockIdx.x] : 0u; + __nv_bfloat16 const* qkv_src = FQKV ? p.packed + (mega_calls & 1u) * kMaxTok * kQkvN : p.packed; + M3FA_PROF(1); + int const par = 0; + if constexpr (FQKV) + { + if (qkv_u >= 0) + { + qkv_mma(p, smem, gbar, tmem_slot, tid, qkv_u); + } + } + + if constexpr (FQKV) + { + if (qkv_u >= 0) + { + if (warp >= kComputeWarps) + { + qkv_epilogue(p, smem, gbar, tmem_slot, tid, qkv_u, par, mega_calls); + } + else + { + // The attention regions overlay the weight ring: every weight byte has gone on into TMEM (or, with + // the MMAs reading the ring, every MMA has completed) first. + mbar_wait_cta(gbar + 16 * kQkvRing + (kQkvASmem ? 8 : 0), 0); + } + } + if (any_valid) + { + dequant_k(); + } +#ifndef M3FA_MEGA_EARLYV +#define M3FA_MEGA_EARLYV 1 // attention CTAs dequantize V here, while the QKV rows are still on their way +#endif +#ifndef M3FA_MEGA_EARLYCLW +#define M3FA_MEGA_EARLYCLW 1 // and pass the cluster barrier (peers' merge mbarriers initialized) here too +#endif + if (M3FA_MEGA_EARLYV && any_valid && tid >= kVStart) + { + dequant_v(); + } + if constexpr (kEarlyClw && !S::kRecvAlias) + { + asm volatile("barrier.cluster.wait.aligned;\n" ::: "memory"); + } + M3FA_PROF(14); + } + + if (any_valid) // CTA-uniform + { + // ---------------- Producer ---------------- + // Q rows of this cluster's tokens in every CTA. The new K/V of every token up to this cluster's last one + // whose position falls in this CTA's blocks go straight into the fp16 tiles; the CTA writes the caches, and + // the index-K, only for its own tokens. + uint2 raw[kMaxT]; + float fq[FQKV ? kMaxT : 1][4]; + bool act[kMaxT]; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + raw[k] = make_uint2(0u, 0u); + act[k] = false; + if (i < kTasks) + { + int t, col; + bool active = true; + if (i >= kHq * QG) + { + int const kv = i - kHq * QG; + t = kv / 3; + int const kind = kv % 3; + col = kind == 0 ? kColK : (kind == 1 ? kColV : kColIdxK); + int const pos = seq_len - QLEN + t; + active = t < tfirst + QG && pos >= 0 && min(pos >> 7, kBlocks - 1) / BPC == cr && slot[t] >= 0 + && slot[t] / kPage < p.num_pages && (kind != 2 || t >= tfirst); + } + else + { + t = tfirst + i / kHq; + col = i % kHq; + } + act[k] = active; + if (active) + { + uint2 const* src = reinterpret_cast( + qkv_src + static_cast(tok0 + t) * p.packed_stride + col * kHeadDim + lane * 4); + if constexpr (FQKV) + { + uint2 v; + bool pending; + do + { + asm volatile("ld.relaxed.gpu.global.v2.b32 {%0, %1}, [%2];\n" + : "=r"(v.x), "=r"(v.y) + : "l"(src) + : "memory"); + pending = (v.x & 0xFFFFu) == 0xFFFFu || (v.x >> 16) == 0xFFFFu || (v.y & 0xFFFFu) == 0xFFFFu + || (v.y >> 16) == 0xFFFFu; + } while (pending); + raw[k] = v; + } + else + { + raw[k] = *src; + } + } + } + } + M3FA_PROF(15); + // While the packed-QKV loads are in flight: dequantize this thread's K chunks to fp16 (exact: e2m1 x e4m3 + // fits fp16). Rows past the longest attended length are zeroed, as the Triton producer masks its loads; the + // new tokens' rows are written by their producer tasks instead. (Building the tiles before the PDL wait was + // measured slower: the cache loads issued at entry are still in flight there, so the wait moves later.) + if constexpr (!FQKV) + { + dequant_k(); + } +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + int const i = warp + kWarps * k; + if (i >= kTasks || !act[k]) + { + continue; + } + float e[4]; + { + float2 const a = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw[k].x)); + float2 const b = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&raw[k].y)); + e[0] = a.x; + e[1] = a.y; + e[2] = b.x; + e[3] = b.y; + } + if (i < kHq * QG) + { + norm_rope(e, qw, p.eps, cs[k], sn[k], lane); + __nv_fp8x2_e4m3 const lo(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const hi(make_float2(e[2], e[3])); + uint2 hv; + hv.x = e4m3x2_to_h2(lo.__x); + hv.y = e4m3x2_to_h2(hi.__x); + *reinterpret_cast(smem + S::kQ + kmaj_off(i, lane * 4)) = hv; // row i = token * 16 + head + continue; + } + int const t = (i - kHq * QG) / 3; + int const kind = (i - kHq * QG) % 3; + int const spage = slot[t] / kPage; + int const within = slot[t] & (kPage - 1); + if (kind == 2) + { + norm_rope(e, ikw, p.eps, cs[k], sn[k], lane); + __nv_fp8x2_e4m3 const lo(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const hi(make_float2(e[2], e[3])); + wr_kind[k] = 2; + wr_val[k] = static_cast(lo.__x) | (static_cast(hi.__x) << 16); + wr_ptr[k] = p.idx_k + static_cast(spage) * p.i_page + static_cast(within) * p.i_tok + + lane * 4; + continue; + } + bool const isV = kind == 1; + if (!isV) + { + norm_rope(e, kw, p.eps, cs[k], sn[k], lane); + } + uint8_t sbyte = 0; + uint16_t const fp4 = quantizeMinimaxM3Fp4x4(e, isV ? vqs : kqs, &sbyte); + if (t >= tfirst) + { + int const scale_col = lane / 4; + int const scale_within = isV ? (within / 4) * 32 + scale_col * 4 + within % 4 : within * 8 + scale_col; + wr_kind[k] = 1; + wr_val[k] = static_cast(fp4) | (static_cast(sbyte) << 16); + wr_ptr[k] = p.kv_data + static_cast(spage) * p.d_page + (isV ? p.d_plane : 0) + within * 64 + + lane * 2; + wr_sptr[k] = p.kv_scale + static_cast(spage) * p.s_page + (isV ? p.s_plane : 0) + scale_within; + } + if (new_row[t] >= 0) + { + // The attention reads this page: its fp16 tile row is the freshly quantized row, dequantized like + // the rest of the tile. + int const h = new_row[t] / kPage; + int const row = new_row[t] % kPage; + uint32_t const sc2 = e4m3_to_h2_splat(sbyte); + uint2 hv; + hv.x = hmul2_u(e2m1x2_to_h2(fp4 & 0xFFu), sc2); + hv.y = hmul2_u(e2m1x2_to_h2(fp4 >> 8), sc2); + *reinterpret_cast( + smem + (isV ? S::kV16 : S::kK16) + h * S::kTileBytes + kmaj_off(row, lane * 4)) + = hv; + } + } + // K tiles, Q rows and patched rows -> tensor core reads. Only threads that wrote fence: the fence also + // waits for the thread's outstanding loads, and the producer warps still have their V chunks in flight. + bool wrote = tid < kComputeThreads; +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + wrote |= act[k]; + } + if (wrote) + { + M3FA_FENCE_PROXY_ASYNC(); + } + M3FA_PROF(13); + __syncthreads(); +#if M3FA_TRIGGER_AT == 2 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Every K/V byte this CTA reads has arrived (the V chunks sit in registers): release the dependents now so + // their weight streams do not queue ahead of the KV reads (the megakernel releases late, from its o_proj). + if constexpr (!FQKV || M3FA_MEGA_TRIGGER == 1) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + tmem = *tmem_slot; + M3FA_PROF(2); + // ---------------- S^T = K Q^T per block on the tensor core: TMEM lane = token, column = query row ---------- +#ifndef M3FA_MEGA_WARPMMA +#define M3FA_MEGA_WARPMMA 1 // attention MMAs issued by warp 0's elected lane with warp-uniform operands (not tid 0) +#endif + if (!kRegCore && (kWarpMma ? warp == 0 : tid == 0)) + { + // Operands computed by the whole warp stay in uniform registers; one elected lane issues. + uint32_t const tm_u = tmem; + bool const issue = kWarpMma ? elect_one() : true; + if (issue) + { +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + uint32_t const a0 = smem_u32(smem + S::kK16 + h * S::kTileBytes); + uint32_t const b0 = smem_u32(smem + S::kQ); +#pragma unroll + for (int ks = 0; ks < kHeadDim / 16; ++ks) + { + umma_f16(tm_u + h * kRows, umma_desc(a0 + ks * 256, 128, 2048), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 0, 0), ks > 0); + } + } + umma_commit(mma_bar); + } + if (kWarpMma) + { + __syncwarp(); + } + } + if (tid >= kVStart && !(FQKV && M3FA_MEGA_EARLYV) && !(!FQKV && M3FA_V15_EARLYV)) + { + // The producer warps dequantize V while QK and the softmax run. + dequant_v(); + } + if constexpr (kRegCore) + { + // ---------------- Register core: warp w of block h = w / kWpb owns 8 * kNt tokens for S and 8 * kNt + // output dims for O; m16n8k16 fp16 MMAs with fp32 accumulators, the softmax in the S fragments ---------- + constexpr int kWpb = kComputeWarps / BPC; // warps per block + constexpr int kNt = 16 / kWpb; // 8-token (S) or 8-dim (O) tiles per warp + float* redx_m = reinterpret_cast(smem + S::kRedX); + float* redx_l = redx_m + BPC * kRows * kWpb; + if (warp < kComputeWarps) + { + int const h = warp / kWpb; + int const wib = warp % kWpb; + int const g = lane >> 2; + int const q4 = lane & 3; + uint32_t const ktile = smem_u32(smem + S::kK16 + h * S::kTileBytes); + uint32_t const vtile = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const qrows = smem_u32(smem + S::kQ); + uint8_t* ptile = smem + S::kK16 + h * S::kTileBytes; // P over this block's K tile, after QK + auto block_bar = [&]() + { + if (h == 0) + { + asm volatile("bar.sync 14, %0;\n" ::"n"(kWpb * 32) : "memory"); + } + else + { + asm volatile("bar.sync 15, %0;\n" ::"n"(kWpb * 32) : "memory"); + } + }; + // S = Q K^T: rows (c, g / g + 8), tokens 8 tb + 2 q4 (+1). + float sv[QG][kNt][4]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + sv[c][nt][e] = 0.0F; + } + } + } +#pragma unroll + for (int ks = 0; ks < kHeadDim / 16; ++ks) + { + uint32_t bk[kNt][2]; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const tb = wib * kNt + nt; + ldsm_x2(ktile + tb * 2048 + (ks * 2 + ((lane >> 3) & 1)) * 128 + (lane & 7) * 16, bk[nt]); + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { + uint32_t aq[4]; + ldsm_x4( + qrows + (c * 2 + ((lane >> 3) & 1)) * 2048 + (ks * 2 + (lane >> 4)) * 128 + (lane & 7) * 16, + aq); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + mma_f16_16816(sv[c][nt], aq, bk[nt]); + } + } + } + M3FA_PROF(7); + // Row maxima (log2 domain, masked as the producer masks), quad then across the block's warps. + float mrow[QG][2]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { + float mx0 = kNegInf, mx1 = kNegInf; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + int const tok = cta_tok0 + h * kPage + (wib * kNt + nt) * 8 + 2 * q4 + (e & 1); + bool const valid = tok < kv_max && tok < kv_len[c]; + float const v = valid ? sv[c][nt][e] * sl2e : kNegInf; + sv[c][nt][e] = v; + if (e < 2) + { + mx0 = fmaxf(mx0, v); + } + else + { + mx1 = fmaxf(mx1, v); + } + } + } + mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, 1)); + mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, 2)); + mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, 1)); + mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, 2)); + if (q4 == 0) + { + redx_m[(h * kRows + c * 16 + g) * kWpb + wib] = mx0; + redx_m[(h * kRows + c * 16 + g + 8) * kWpb + wib] = mx1; + } + } + block_bar(); // every warp of the block is past QK (P may overwrite the K tile) and wrote its maxima + M3FA_PROF(8); +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int half = 0; half < 2; ++half) + { + float const* rm = redx_m + (h * kRows + c * 16 + g + 8 * half) * kWpb; + float m = kNegInf; +#pragma unroll + for (int w4 = 0; w4 < kWpb; w4 += 4) + { + float4 const v = *reinterpret_cast(rm + w4); + m = fmaxf(m, fmaxf(fmaxf(v.x, v.y), fmaxf(v.z, v.w))); + } + mrow[c][half] = m; + } + // P = exp2(s - m) as fp16 into [row][token] core matrices; the fp32 values feed the row sums. + float l0 = 0.0F, l1 = 0.0F; +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + float a[4]; +#pragma unroll + for (int e = 0; e < 4; ++e) + { + float const m = mrow[c][e >> 1]; + a[e] = m == kNegInf ? 0.0F : ex2(sv[c][nt][e] - m); + } + l0 += a[0] + a[1]; + l1 += a[2] + a[3]; + int const t0 = (wib * kNt + nt) * 8 + 2 * q4; + int const r0 = c * 16 + g; + __half2 const p0 = __floats2half2_rn(a[0], a[1]); + __half2 const p1 = __floats2half2_rn(a[2], a[3]); + *reinterpret_cast<__half2*>( + ptile + (r0 >> 3) * 2048 + (t0 >> 3) * 128 + (r0 & 7) * 16 + (t0 & 7) * 2) + = p0; + *reinterpret_cast<__half2*>( + ptile + ((r0 + 8) >> 3) * 2048 + (t0 >> 3) * 128 + (r0 & 7) * 16 + (t0 & 7) * 2) + = p1; + } + l0 += __shfl_xor_sync(0xffffffffu, l0, 1); + l0 += __shfl_xor_sync(0xffffffffu, l0, 2); + l1 += __shfl_xor_sync(0xffffffffu, l1, 1); + l1 += __shfl_xor_sync(0xffffffffu, l1, 2); + if (q4 == 0) + { + redx_l[(h * kRows + c * 16 + g) * kWpb + wib] = l0; + redx_l[(h * kRows + c * 16 + g + 8) * kWpb + wib] = l1; + } + } + M3FA_PROF_T(19, 0); + } + __syncthreads(); // P, the V tiles (producer warps) and the row sums visible + M3FA_PROF(9); + if (warp < kComputeWarps) + { + int const h = warp / kWpb; + int const wib = warp % kWpb; + int const g = lane >> 2; + int const q4 = lane & 3; + uint32_t const vtile = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const ptile = smem_u32(smem + S::kK16 + h * S::kTileBytes); + float* obuf = reinterpret_cast(smem + S::kK16 + h * S::kTileBytes + S::kObufOff); + // O = P V: rows (c, g / g + 8), dims 8 dt + 2 q4 (+1). + float ov[QG][kNt][4]; +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + ov[c][nt][e] = 0.0F; + } + } + } +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + uint32_t bv[kNt][2]; + int const tok = ks * 16 + ((lane >> 3) & 1) * 8 + (lane & 7); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const dt = wib * kNt + nt; + ldsm_x2_trans(vtile + (tok >> 3) * 2048 + dt * 128 + (tok & 7) * 16, bv[nt]); + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { + uint32_t ap[4]; + ldsm_x4( + ptile + (c * 2 + ((lane >> 3) & 1)) * 2048 + (ks * 2 + (lane >> 4)) * 128 + (lane & 7) * 16, + ap); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + mma_f16_16816(ov[c][nt], ap, bv[nt]); + } + } + } +#pragma unroll + for (int c = 0; c < QG; ++c) + { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) + { + int const d = (wib * kNt + nt) * 8 + 2 * q4; + int const r0 = c * 16 + g; + *reinterpret_cast(obuf + r0 * kOPitch + d) = make_float2(ov[c][nt][0], ov[c][nt][1]); + *reinterpret_cast(obuf + (r0 + 8) * kOPitch + d) + = make_float2(ov[c][nt][2], ov[c][nt][3]); + } + } + } + else if (tid - kComputeThreads < BPC * kRows) + { + // Row statistics of every (block, row): LSE and the PV scale, as the producer derives them. + int const r = tid - kComputeThreads; + float const* rm = redx_m + r * kWpb; + float const* rl = redx_l + r * kWpb; + float m = kNegInf; + float l = 0.0F; +#pragma unroll + for (int w4 = 0; w4 < kWpb; w4 += 4) + { + float4 const vm = *reinterpret_cast(rm + w4); + float4 const vl = *reinterpret_cast(rl + w4); + m = fmaxf(m, fmaxf(fmaxf(vm.x, vm.y), fmaxf(vm.z, vm.w))); + l += (vl.x + vl.y) + (vl.z + vl.w); + } + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[r] = scale * vgs; + row_lse[r] = lse; + } + __syncthreads(); // O partials and row statistics visible to the push + } + else + { + if constexpr (kRs2) + { + // ---------------- Split softmax: warps 0..3 rows 0..7, warps 4..7 rows 8..15 (thread = block token) + // -------- + if (warp < 8) + { + int const rg = warp / 4; + int const wq = warp % 4; + int const tk = wq * 32 + lane; + auto rg_bar = [&]() + { + if (rg == 0) + { + asm volatile("bar.sync 10, 128;\n" ::: "memory"); + } + else + { + asm volatile("bar.sync 11, 128;\n" ::: "memory"); + } + }; + int const rs8 = ((lane >> 4) & 1) * 4 + ((lane >> 3) & 1) * 2 + ((lane >> 2) & 1); + mbar_wait_cta(mma_bar, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF(7); + float v[8]; + tmem_ld8(tmem + ((wq * 32) << 16) + 8 * rg, v); + bool const valid = cta_tok0 + tk < kv_max && cta_tok0 + tk < kv_len[0]; + float a[8]; +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] = valid ? v[j2] * sl2e : kNegInf; + a[j2] = v[j2]; + } + float const mw = warp_reduce_scatter8(a, lane); + if ((lane & 3) == 0) + { + red_m[(8 * rg + rs8) * 4 + wq] = mw; + } + rg_bar(); + M3FA_PROF(8); + uint8_t* pdst = smem + S::kP + (tk >> 3) * 128 + (tk & 7) * 16; + float mrow = kNegInf; + if (lane < 8) + { + float4 const mq = *reinterpret_cast(red_m + (8 * rg + lane) * 4); + mrow = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + } + float pa[8]; +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + float const m = __shfl_sync(0xffffffffu, mrow, j2); + pa[j2] = m == kNegInf ? 0.0F : ex2(v[j2] - m); + } + if (tid == 0) + { + M3FA_CLK(41); + } + uint32_t pw[4]; +#pragma unroll + for (int j2 = 0; j2 < 4; ++j2) + { + __half2 const hp = __floats2half2_rn(pa[2 * j2], pa[2 * j2 + 1]); + pw[j2] = *reinterpret_cast(&hp); + } + *reinterpret_cast(pdst + rg * 2048) = make_uint4(pw[0], pw[1], pw[2], pw[3]); + M3FA_PROF_T(18, 0); + if (tid == 0) + { + M3FA_CLK(42); + } + M3FA_FENCE_PROXY_ASYNC(); // P -> tensor core reads + if (tid == 0) + { + M3FA_CLK(43); + } + if (warp == 0) + { + asm volatile("bar.sync 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + if (tid == 0) + { + M3FA_CLK(44); + } + } + else + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + if (kWarpMma ? warp == 0 : tid == 0) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm_u = tmem; + uint32_t const a0 = smem_u32(smem + S::kV16); + uint32_t const b0 = smem_u32(smem + S::kP); + if (kWarpMma ? elect_one() : true) + { +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + umma_f16(tm_u + BPC * kRows, umma_desc(a0 + ks * 4096, 2048, 128), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 1, 1), ks > 0); + } + umma_commit(mma_bar); + } + if (kWarpMma) + { + __syncwarp(); + } + } + M3FA_PROF_T(19, 0); + float const lw = warp_reduce_scatter8(pa, lane); + if ((lane & 3) == 0) + { + red_l[(8 * rg + rs8) * 4 + wq] = lw; + } + rg_bar(); + M3FA_PROF_T(20, 0); + if (wq == 0 && lane < 8) + { + int const r = 8 * rg + lane; + float4 const mq = *reinterpret_cast(red_m + r * 4); + float4 const lq = *reinterpret_cast(red_l + r * 4); + float const m = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + float const l = (lq.x + lq.y) + (lq.z + lq.w); + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[r] = scale * vgs; + row_lse[r] = lse; + } + } + __syncthreads(); // row scales and LSEs visible (PV was issued above) + M3FA_PROF(9); + if (warp < 8) + { + mbar_wait_cta(mma_bar, 1); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + } + else + { + // ---------------- Softmax: thread = block token; the (block, 16-row chunk) pairs are spread over the + // four groups of four compute warps (each group covers the 128 TMEM lanes = block tokens) + // ---------------- + if (warp < kComputeWarps && warp / 4 < kPairs) + { + int const g = warp / 4; + int const wq = warp % 4; + int const tk = wq * 32 + lane; + auto group_bar = [&]() + { + if (g == 0) + { + asm volatile("bar.sync 10, 128;\n" ::: "memory"); + } + else if (g == 1) + { + asm volatile("bar.sync 11, 128;\n" ::: "memory"); + } + else if (g == 2) + { + asm volatile("bar.sync 12, 128;\n" ::: "memory"); + } + else + { + asm volatile("bar.sync 13, 128;\n" ::: "memory"); + } + }; + // Row r(lane) this lane holds after the reduce-scatter: bits 4..1 of the lane. + int const rs_row + = ((lane >> 4) & 1) * 8 + ((lane >> 3) & 1) * 4 + ((lane >> 2) & 1) * 2 + ((lane >> 1) & 1); + mbar_wait_cta(mma_bar, 0); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + M3FA_PROF(7); + float v[kPairsPerGroup][16]; +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + tmem_ld16(tmem + ((wq * 32) << 16) + h * kRows + c * 16, v[k]); + bool const valid + = cta_tok0 + h * kPage + tk < kv_max && cta_tok0 + h * kPage + tk < kv_len[c]; + float a[16]; +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[k][j] = valid ? v[k][j] * sl2e : kNegInf; + a[j] = v[k][j]; + } + float const mw = warp_reduce_scatter16(a, lane); + if ((lane & 1) == 0) + { + red_m[((h * QG + c) * 16 + rs_row) * 4 + wq] = mw; + } + } + } + group_bar(); + M3FA_PROF(8); + float pa[kPairsPerGroup][16]; +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + uint8_t* pdst = smem + S::kP + h * kRows * kPage * 2 + (tk >> 3) * 128 + (tk & 7) * 16; + // Row maxima: lane j (< 16) reduces row j's four warp partials, then broadcasts it. + float mrow = kNegInf; + if (lane < 16) + { + float4 const mq + = *reinterpret_cast(red_m + ((h * QG + c) * 16 + lane) * 4); + mrow = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + } +#pragma unroll + for (int j = 0; j < 16; ++j) + { + float const m = __shfl_sync(0xffffffffu, mrow, j); + pa[k][j] = m == kNegInf ? 0.0F : ex2(v[k][j] - m); + } + if (tid == 0) + { + M3FA_CLK(41); + } + uint32_t pw[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) + { + __half2 const hp = __floats2half2_rn(pa[k][2 * j], pa[k][2 * j + 1]); + pw[j] = *reinterpret_cast(&hp); + } + // P^T as the MN-major B operand: 8 query rows of one token are one 16-byte core-matrix row. + *reinterpret_cast(pdst + (2 * c) * 2048) = make_uint4(pw[0], pw[1], pw[2], pw[3]); + *reinterpret_cast(pdst + (2 * c + 1) * 2048) + = make_uint4(pw[4], pw[5], pw[6], pw[7]); + M3FA_PROF_T(18, 0); + } + } + if (tid == 0) + { + M3FA_CLK(42); + } + M3FA_FENCE_PROXY_ASYNC(); // P -> tensor core reads + if (tid == 0) + { + M3FA_CLK(43); + } + // PV starts as soon as P and the V tiles are in: thread 0 waits for the softmax warps and the + // producer warps (V); the row sums and statistics below overlap the MMAs. + if (warp == 0) + { + asm volatile("bar.sync 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + else + { + asm volatile("bar.arrive 6, %0;\n" ::"n"(kPvBarThreads) : "memory"); + } + if (kWarpMma ? warp == 0 : tid == 0) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + uint32_t const tm_u = tmem; + if (kWarpMma ? elect_one() : true) + { +#pragma unroll + for (int h = 0; h < BPC; ++h) + { + uint32_t const a0 = smem_u32(smem + S::kV16 + h * S::kTileBytes); + uint32_t const b0 = smem_u32(smem + S::kP + h * kRows * kPage * 2); +#pragma unroll + for (int ks = 0; ks < kPage / 16; ++ks) + { + umma_f16(tm_u + (BPC + h) * kRows, umma_desc(a0 + ks * 4096, 2048, 128), + umma_desc(b0 + ks * 256, 128, 2048), umma_idesc(kRows, 1, 1), ks > 0); + } + } + umma_commit(mma_bar); + } + if (kWarpMma) + { + __syncwarp(); + } + } + M3FA_PROF_T(19, 0); +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + float const lw = warp_reduce_scatter16(pa[k], lane); + if ((lane & 1) == 0) + { + red_l[((h * QG + c) * 16 + rs_row) * 4 + wq] = lw; + } + } + } + group_bar(); + M3FA_PROF_T(20, 0); + if (wq == 0 && lane < 16) + { +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + int const r = (h * QG + c) * 16 + lane; + float4 const mq = *reinterpret_cast(red_m + r * 4); + float4 const lq = *reinterpret_cast(red_l + r * 4); + float const m = fmaxf(fmaxf(mq.x, mq.y), fmaxf(mq.z, mq.w)); + float const l = (lq.x + lq.y) + (lq.z + lq.w); + float const lse = l > 0.0F ? m + lg2(l) : kNegInf; + float const scale = lse > kNegInf ? ex2(m - lse) : 0.0F; + row_scale[h * kRows + c * 16 + lane] = scale * vgs; + row_lse[h * kRows + c * 16 + lane] = lse; + } + } + } + } + __syncthreads(); // row scales and LSEs visible (PV was issued above) + M3FA_PROF(9); + if (warp < kComputeWarps && warp / 4 < kPairs) + { + mbar_wait_cta(mma_bar, 1); + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + } + } + } + M3FA_PROF(3); + } + else + { + asm volatile("cp.async.wait_all;\n" ::: "memory"); +#if M3FA_TRIGGER_AT == 2 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (!FQKV || M3FA_MEGA_TRIGGER == 1) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif + for (int r = tid; r < BPC * kRows; r += kThreads) + { + row_lse[r] = kNegInf; + row_scale[r] = 0.0F; + } + __syncthreads(); + } + + // ---------------- Push every partial row to the CTA that merges its head (CTA r merges heads r*BPC..) -------- + if constexpr (S::kRecvAlias) + { + // Peers' partials overlay this CTA's tiles: every CTA's PV MMAs have completed. + asm volatile("barrier.cluster.arrive.aligned;\n" ::: "memory"); + } + if (!(kEarlyClw && !S::kRecvAlias)) + { + asm volatile("barrier.cluster.wait.aligned;\n" ::: "memory"); // peers' mbarriers are initialized + } + M3FA_PROF(4); + uint32_t const recv_base = smem_u32(recv); + if (kRs2 && warp < 8) + { + // Split push: warp group rg pushes rows 8 rg .. 8 rg + 7 (row j goes to CTA j, which merges head j). + int const rg = warp / 4; + int const d = (warp % 4) * 32 + lane; + float v[8]; + if (any_valid) + { + tmem_ld8(tmem + ((warp % 4) * 32 << 16) + BPC * kRows + 8 * rg, v); + } + else + { +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] = 0.0F; + } + } +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + v[j2] *= row_scale[8 * rg + j2]; + } +#pragma unroll + for (int j2 = 0; j2 < 8; ++j2) + { + uint32_t const dst_cta = 8 * rg + j2; + uint32_t const row_addr = recv_base + static_cast((cr * kOPitch + d) * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr, dst_cta)), + "r"(__float_as_uint(v[j2])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + else if (!kRs2 && warp < kComputeWarps && warp / 4 < kPairs) + { + int const g = warp / 4; + int const d = (warp % 4) * 32 + lane; // output dim +#pragma unroll + for (int k = 0; k < kPairsPerGroup; ++k) + { + int const q = g + 4 * k; + if (q < kPairs) + { + int const h = q % BPC; + int const c = q / BPC; + int const chunk = cr * BPC + h; + float v[16]; + if (any_valid && kRegCore) + { + float const* obuf + = reinterpret_cast(smem + S::kK16 + h * S::kTileBytes + S::kObufOff); +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] = obuf[(c * 16 + j) * kOPitch + d]; + } + } + else if (any_valid) + { + tmem_ld16(tmem + ((warp % 4) * 32 << 16) + (BPC + h) * kRows + c * 16, v); + } + else + { +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] = 0.0F; + } + } +#pragma unroll + for (int j = 0; j < 16; ++j) + { + v[j] *= row_scale[h * kRows + c * 16 + j]; + } + // One scalar store per row: a vector store's lanes would address several destination CTAs at once, + // which measured several times slower. +#pragma unroll + for (int j = 0; j < 16; ++j) + { + uint32_t const dst_cta = j / BPC; + uint32_t const row_addr + = recv_base + static_cast((((chunk * QG + c) * BPC + j % BPC) * kOPitch + d) * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr, dst_cta)), + "r"(__float_as_uint(v[j])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + } + } + else if (tid >= kComputeThreads) + { + // Every partial row's LSE (BPC * kRows of them; more than the producer threads for the widest geometries). + for (int r = tid - kComputeThreads; r < BPC * kRows; r += kThreads - kComputeThreads) + { + int const h = r / kRows; + int const c = (r % kRows) / 16; + int const j = r % 16; + uint32_t const dst_cta = j / BPC; + uint32_t const row_addr + = recv_base + static_cast((((cr * BPC + h) * QG + c) * BPC + j % BPC) * kOPitch * 4); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_shared(row_addr + kHeadDim * 4, dst_cta)), + "r"(__float_as_uint(row_lse[r])), "r"(mapa_shared(mbar, dst_cta)) + : "memory"); + } + } + M3FA_PROF(10); +#pragma unroll + for (int k = 0; k < kMaxT; ++k) + { + if (wr_kind[k] == 1) + { + *reinterpret_cast(wr_ptr[k]) = static_cast(wr_val[k] & 0xFFFFu); + if (lane % 4 == 0) + { + *wr_sptr[k] = static_cast(wr_val[k] >> 16); + } + } + else if (wr_kind[k] == 2) + { + *reinterpret_cast(wr_ptr[k]) = wr_val[k]; + } + } + { + uint32_t ready = 0; +#ifdef M3FA_DBG_HANG + long long spins = 0; +#endif + while (!ready) + { +#ifdef M3FA_DBG_HANG + if (++spins == (1LL << 24)) + { + printf("M3FA hang: block %d thread %d merge recv\n", blockIdx.x, threadIdx.x); + } +#endif + 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}\n" + : "=r"(ready) + : "r"(mbar) + : "memory"); + } + } + M3FA_PROF(5); + + // ---------------- Merge heads [cr * BPC, cr * BPC + BPC) (split-K merge kernel's LSE weighting) ---------------- + constexpr int kMergeRows = QG * BPC; // (token, local head) pairs, 16 chunk weights each + // Wide geometries: chunk weights once per merged row (into the softmax scratch, free by now), then the weighted + // sums. Narrow ones recompute the weights per output element: no extra barrier on the critical path. + constexpr bool kRowWeights = kMergeRows >= 4; + float* wmerge = red_m; // [kMergeRows][kBlocks]; red_m holds BPC * kRows * 4 = kMergeRows * 64 floats + [[maybe_unused]] auto chunk_weights = [&](int mr, float(&lc)[kBlocks]) + { + float lmax = kNegInf; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + lc[c] = recv[(c * kMergeRows + mr) * kOPitch + kHeadDim]; + lmax = fmaxf(lmax, lc[c]); + } + lmax = lmax == kNegInf ? 0.0F : lmax; + float denom = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + lc[c] = ex2(lc[c] - lmax); // now the chunk weight; 0 for an empty chunk + denom += lc[c]; + } + denom = denom > 0.0F ? denom : 1.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + // -0.0 marks an empty chunk (skipped below, as the merge kernel skips zero weights). + lc[c] = lc[c] != 0.0F ? __fdividef(lc[c], denom) : -0.0F; + } + }; + if constexpr (kRowWeights) + { + for (int mr = tid; mr < kMergeRows; mr += kThreads) + { + float lc[kBlocks]; + chunk_weights(mr, lc); +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wmerge[mr * kBlocks + c] = lc[c]; + } + } + __syncthreads(); + } +#pragma unroll + for (int idx = tid; idx < kMergeRows * kHeadDim; idx += kThreads) + { + int const mr = idx >> 7; + int const d = idx & (kHeadDim - 1); + float wn[kBlocks]; + if constexpr (kRowWeights) + { +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wn[c] = wmerge[mr * kBlocks + c]; + } + } + else + { +#ifndef M3FA_MEGA_MERGE_LOCAL +#define M3FA_MEGA_MERGE_LOCAL \ + 0 // 1 (measured +0.6 us slower): every merging thread derives the 16 chunk weights itself (broadcast shared loads, + // independent exps): no shuffle chain; the same arithmetic as chunk_weights +#endif + if constexpr (M3FA_MEGA_MERGE_LOCAL != 0) + { + chunk_weights(mr, wn); + } + else + { + // Narrow geometries: the warp (32 dims of one merged row) derives the 16 chunk weights together, lane c + // holding chunk c's, and broadcasts them. + float lc = (lane < kBlocks) ? recv[(lane * kMergeRows + mr) * kOPitch + kHeadDim] : kNegInf; + float lmax = lc; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + lmax = fmaxf(lmax, __shfl_xor_sync(0xffffffffu, lmax, o)); + } + lmax = lmax == kNegInf ? 0.0F : lmax; + float const wc = lane < kBlocks ? ex2(lc - lmax) : 0.0F; + // The denominator in chunk order from 0.0, as chunk_weights (and the merge kernel) add it: a butterfly + // reduction rounds differently in the last bit. + float denom = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + denom += __shfl_sync(0xffffffffu, wc, c); + } + denom = denom > 0.0F ? denom : 1.0F; + float const wl = wc != 0.0F ? __fdividef(wc, denom) : -0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + wn[c] = __shfl_sync(0xffffffffu, wl, c); + } + } + } + if (tid == 0) + { + M3FA_CLK(45); + } + float acc = 0.0F; +#pragma unroll + for (int c = 0; c < kBlocks; ++c) + { + float const prod = recv[(c * kMergeRows + mr) * kOPitch + d] * wn[c]; + acc = __float_as_uint(wn[c]) != 0x80000000u ? acc + prod : acc; + } + if (tid == 0) + { + M3FA_CLK(46); + } + int const tl = mr / BPC; + int const head = cr * BPC + mr % BPC; + __nv_bfloat16 const ob = __float2bfloat16_rn(acc); + int const out_row = tok0 + tfirst + tl; + p.out[static_cast(out_row) * p.out_stride + head * kHeadDim + d] = ob; + if constexpr (MXFP8) + { + float const x = __bfloat162float(ob); + float amax = fabsf(x); +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, mask)); + } + float const sfv = amax * rcp_approx_ftz(448.0F); + int const bits = __float_as_int(sfv); + int ex = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0 ? 1 : 0); + ex = min(ex, 254); + float scale = __int_as_float((254 - ex) << 23); + scale = amax != 0.0F ? scale : 0.0F; + __nv_fp8_e4m3 const q8v(x * scale); + int64_t const q8off = static_cast(out_row) * (p.num_sf_k_tiles * 128) + head * kHeadDim + d; + if constexpr (FQKV) + { + // Megakernel hand-off to the o_proj units (Lamport): this call's parity buffer, the other one reset to + // 0xFF. 0xFF is never written: a NaN code is canonicalized to 0x7F; the UE8M0 code is at most 254. + uint8_t const code = q8v.__x == 0xFFu ? static_cast(0x7Fu) : q8v.__x; + int64_t const cur = (mega_calls & 1u) * kMaxTok * kOpK; + int64_t const nxt = ((mega_calls & 1u) ^ 1u) * kMaxTok * kOpK; + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.q8 + cur + q8off), "h"(static_cast(code)) + : "memory"); + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.q8 + nxt + q8off), "h"(static_cast(0xFFu)) + : "memory"); + } + else + { + p.q8[q8off] = q8v.__x; + } + if (lane == 0) + { + int const k_sf = head * (kHeadDim / 32) + (d >> 5); + int64_t const off = static_cast(out_row / 128) * p.num_sf_k_tiles * 512 + (k_sf / 4) * 512 + + (out_row % 32) * 16 + ((out_row % 128) / 32) * 4 + k_sf % 4; + if constexpr (FQKV) + { + int64_t const cur = (mega_calls & 1u) * kOpSfBytes; + int64_t const nxt = ((mega_calls & 1u) ^ 1u) * kOpSfBytes; + asm volatile( + "st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.sf + cur + off), "h"(static_cast(ex)) + : "memory"); + asm volatile("st.relaxed.gpu.global.u8 [%0], %1;\n" ::"l"(p.sf + nxt + off), + "h"(static_cast(0xFFu)) + : "memory"); + } + else + { + p.sf[off] = static_cast(ex); + } + } + } + } + M3FA_PROF(60); + if (any_valid) + { + // Every TMEM read completed above (tcgen05.wait::ld); release the columns. + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); + } + __syncthreads(); + if (FQKV && tid == 0) + { + p.flags[32 + blockIdx.x] = mega_calls + 1u; + } + M3FA_PROF(61); + if (warp == kWarps - 1) + { + asm volatile("tcgen05.fence::after_thread_sync;\n" ::: "memory"); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" ::"r"(*tmem_slot), "r"(S::kTmemCols) + : "memory"); + } + M3FA_PROF(6); +#else + __trap(); +#endif // MINIMAX_M3_ATTN_DEVICE_ENABLED +} + +template +cudaError_t launch_impl(Params const& prm, int num_reqs, bool pdl, cudaStream_t stream, int* max_clusters) +{ + auto kernel = m3_fused_decode_attn_kernel; + constexpr int smem = SmemX::kBytes; + constexpr int kCs = kBlocks / BPC; + static bool const configured = [&]() + { + bool ok = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem) == cudaSuccess; + ok = ok + && (kCs <= 8 + || cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) == cudaSuccess); + if (!ok) + { + cudaGetLastError(); + } + return ok; + }(); + if (!configured) + { + return cudaErrorInvalidConfiguration; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(FQKV ? kQkvGrid : kCs * (QLEN / QG) * num_reqs); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cudaLaunchAttribute attr[2]; + attr[0].id = cudaLaunchAttributeClusterDimension; + attr[0].val.clusterDim.x = kCs; + attr[0].val.clusterDim.y = 1; + attr[0].val.clusterDim.z = 1; + attr[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[1].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + if (max_clusters != nullptr) + { + cfg.numAttrs = 1; + cudaError_t const err = cudaOccupancyMaxActiveClusters(max_clusters, kernel, &cfg); + if (err != cudaSuccess) + { + cudaGetLastError(); + } + return err; + } + cfg.numAttrs = 2; + return cudaLaunchKernelEx(&cfg, kernel, prm); +} + +// The (QLEN, QG, BPC) geometries the target launches: default_geometry() of every step shape it routes here (one +// query token per 16-CTA cluster, or KV blocks paired into 8-CTA clusters), always with the MXFP8 hand-off. +#define M3_FA_GEOMETRIES(X) \ + X(1, 1, 1) \ + X(1, 1, 2) X(2, 1, 1) X(2, 1, 2) X(3, 1, 1) X(4, 1, 1) X(5, 1, 1) X(6, 1, 1) X(6, 1, 2) X(7, 1, 1) X(8, 1, 1) + +cudaError_t dispatch(Params const& prm, int qlen, int qg, int bpc, int num_reqs, bool pdl, cudaStream_t stream) +{ +#define M3_FA_CASE(Q, G, B) \ + if (qlen == Q && qg == G && bpc == B) \ + return launch_impl(prm, num_reqs, pdl, stream, nullptr); + M3_FA_GEOMETRIES(M3_FA_CASE) +#undef M3_FA_CASE + return cudaErrorInvalidValue; +} + +} // namespace + +bool minimaxM3FusedDecodeAttnSupports(int qlen, int qg, int bpc) +{ +#define M3_FA_CASE(Q, G, B) \ + if (qlen == Q && qg == G && bpc == B) \ + return true; + M3_FA_GEOMETRIES(M3_FA_CASE) +#undef M3_FA_CASE + return false; +} + +cudaError_t invokeMinimaxM3FusedDecodeAttn( + MinimaxM3FusedDecodeAttnParams const& params, int qlen, int qg, int bpc, int numReqs, bool pdl, cudaStream_t stream) +{ + return dispatch(params, qlen, qg, bpc, numReqs, pdl, stream); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.h b/cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.h new file mode 100644 index 000000000000..663c2d0aad91 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.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 +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +//! Kernel parameters of the fused decode attention (the device code's Params). +struct MinimaxM3FusedDecodeAttnParams +{ + __nv_bfloat16 const* packed; + int64_t packed_stride; + uint8_t* kv_data; + int64_t d_page, d_plane; + uint8_t* kv_scale; + int64_t s_page, s_plane; + uint8_t* idx_k; + int64_t i_page, i_tok; + int64_t num_pages; + int const* out_cache_loc; + int const* position_ids; + int const* block_table; + int64_t bt_stride; + int bt_cols; + int const* seq_lens; + float const* kv_quant_scale; + float const* k_gs; + float const* v_gs; + __nv_bfloat16 const* q_w; + __nv_bfloat16 const* k_w; + __nv_bfloat16 const* ik_w; + float const* rope; + float eps; + float sm_scale; + __nv_bfloat16* out; + int64_t out_stride; + uint8_t* q8; + uint8_t* sf; + int num_sf_k_tiles; + unsigned long long* prof; // optional per-CTA clock64 probes (-DM3FA_PROFILE builds only) + // Fused QKV projection (FQKV kernels only): out[M, 2560] = mxfp8(x)[M, 6144] @ W^T. + __nv_bfloat16 const* x; // [M, 6144] layer input (the all-reduce output) + int ldx; + int num_tokens; // M + uint8_t const* wsf; // UE8M0 weight scales, 128x4-swizzled + float* part; // [2 parities][splits][8 tokens][2560] partial sums, sentinel at rest + int* counters; // [0]: call parity, [1]: CTAs done with this call (zero at rest) + int attn_clusters; // leading clusters that run the attention + uint8_t const* xq; // mega: the layer input's MXFP8 values [M, 6144] (row stride ldxq) ... + uint8_t const* xsf_act; // ... and its 128x4-swizzled UE8M0 scale atoms (atom k at byte 512 k) + int ldxq; + __nv_bfloat16* qkv_out; // mega: the packed QKV rows [2 parities][8][2560] the attention reads + uint8_t const* wsf_o; // mega: the o_proj weight's scale atoms + __nv_bfloat16* y_out; // mega: the o_proj output [M, 6144] + int* ar_ws; // mega, ar_mode: the next m3_mc_allreduce's multicast exchange (int32 words)... + int const* ar_flags; // ... its call flags (parity of a (token, column slice) slot = flags[t * 8 + slice] & 1) + int ar_rank; + int ar_world; + int ar_maxt; + int ar_mode; // 0: none; 2: push the o_proj rows into the exchange (as m3_mx_gemv's ar mode 2) + unsigned int* flags; // mega: [0, 20) tile counters (4 per call), [32 + CTA] the attention CTAs' targets + long long pf_a0, pf_b0, pf_a1, pf_b1; // mega: L2 prefetch regions issued in the attention window (bytes 0: none) + int pf_hint; // 1: evict_last policy on the prefetch + uint8_t const* wpre; // mega: optional pre-tiled QKV weight [N/128][K/128][16 KB swizzled smem image], or null + alignas(64) CUtensorMap wmap; // the [2560, 6144] E4M3 weight, 128 x 128 boxes, 128-byte swizzle + alignas(64) CUtensorMap wmap_o; // mega: the [6144, 2048] E4M3 o_proj weight, same boxes +}; + +//! Whether the (decode_query_len, query tokens per cluster, KV blocks per CTA) geometry is compiled. The MXFP8 +//! hand-off (q8 / sf) is part of every compiled variant. +bool minimaxM3FusedDecodeAttnSupports(int qlen, int qg, int bpc); + +//! Launch the fused norm/RoPE/KV-insert + identity-selection sparse decode + split-K merge kernel on `stream`. +//! `params` holds the packed QKV rows, the paged caches and the outputs (see the field comments). Returns +//! cudaErrorInvalidValue for a geometry minimaxM3FusedDecodeAttnSupports() rejects. +cudaError_t invokeMinimaxM3FusedDecodeAttn(MinimaxM3FusedDecodeAttnParams const& params, int qlen, int qg, int bpc, + int numReqs, bool pdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/mcAllreduceNorm.cu b/cpp/tensorrt_llm/kernels/minimaxM3/mcAllreduceNorm.cu new file mode 100644 index 000000000000..fc60fb3035a9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/mcAllreduceNorm.cu @@ -0,0 +1,563 @@ +/* + * 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/minimaxM3/mcAllreduceNorm.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmException.h" + +#include +#include +#include +#include + +// The kernels use sm_100-family features (thread-block clusters, cp.async.bulk, st.async, mbarrier transaction +// counts, block-scaled conversions). Every other architecture of the build compiles them to empty bodies; the host +// entry points refuse to launch on those devices. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030) \ + && (defined(__CUDA_ARCH_FEAT_SM100_ALL) || defined(__CUDA_ARCH_FEAT_SM103_ALL) \ + || defined(__CUDA_ARCH_FAMILY_SPECIFIC__)) +#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 + +__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); +} + +__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. +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* __restrict__ ws_uc, int* __restrict__ 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) +{ +#if MINIMAX_M3_KERNELS_ENABLED + __shared__ float warp_sq[THREADS / 32]; + __shared__ alignas(16) float cluster_sq[CLUSTER]; + __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]); + } + } + } + __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); + 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 & 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 (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 trtllm::mxfp8_quantize(normed, True): 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) + { + __nv_fp8_e8m0 code; + code.__x = __nv_cvt_float_to_e8m0(amax * rcp_approx_ftz(448.0f), __NV_SATFINITE, cudaRoundPosInf); + float const out_scale = amax != 0.f ? rcp_approx_ftz(static_cast(code)) : 0.0f; + uint32_t packed[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) + { + __nv_fp8x2_e4m3 const lo(make_float2(s[4 * h] * out_scale, s[4 * h + 1] * out_scale)); + __nv_fp8x2_e4m3 const hi(make_float2(s[4 * h + 2] * out_scale, s[4 * h + 3] * out_scale)); + packed[h] = static_cast(lo.__x) | (static_cast(hi.__x) << 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] + = code.__x; + } + } + } + if constexpr ((OPTS & 4) == 0) + { + __syncthreads(); + if (threadIdx.x == 0) + flags[blockIdx.x] = parity ^ 1; + } +#endif +} + +template +void launch(MinimaxM3McAllreduceNormParams 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 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; + 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); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +constexpr int kCluster = 8; +constexpr int kThreads = 128; + +} // namespace + +bool isMinimaxM3McAllreduceNormSupported(int32_t maxTokens, int32_t opts) +{ + return (maxTokens == 8 && (opts == 7 || opts == 15)) || (maxTokens == 16 && opts == 15) + || (maxTokens == 48 && opts == 15); +} + +void invokeMinimaxM3McAllreduceNorm( + MinimaxM3McAllreduceNormParams const& params, int32_t maxTokens, int32_t opts, bool pdl, cudaStream_t stream) +{ + if (maxTokens == 8 && opts == 7) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 15) + { + 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/minimaxM3/mcAllreduceNorm.h b/cpp/tensorrt_llm/kernels/minimaxM3/mcAllreduceNorm.h new file mode 100644 index 000000000000..b80a78c9e173 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/mcAllreduceNorm.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 + +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 MinimaxM3McAllreduceNormParams +{ + __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}; +}; + +//! The (bucket, epilogue options) pairs this launcher is built for: (8, 7), (8, 15), (16, 15) and (48, 15). +bool isMinimaxM3McAllreduceNormSupported(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 invokeMinimaxM3McAllreduceNorm( + MinimaxM3McAllreduceNormParams 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/minimaxM3/msaKvUpdate.cu b/cpp/tensorrt_llm/kernels/minimaxM3/msaKvUpdate.cu new file mode 100644 index 000000000000..2fb9f768c1bd --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/msaKvUpdate.cu @@ -0,0 +1,70 @@ +/* + * 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/minimaxM3/msaKvUpdate.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ +__global__ void __launch_bounds__(256) + msa_kv_update_kernel(int const* __restrict__ kv_lens, int const* __restrict__ staged, int const* __restrict__ qo, + int const* __restrict__ q_batch_row, int const* __restrict__ q_intra, int const* __restrict__ table, + long long table_ld, int ncols, int page, int batch, int total_q, int* __restrict__ seq_lens_out, + int* __restrict__ cu_kv, int* __restrict__ out_loc, int* __restrict__ n_valid) +{ + for (int b = threadIdx.x; b < batch; b += blockDim.x) + seq_lens_out[b] = min(kv_lens[b], staged[b]); + if (cu_kv != nullptr && threadIdx.x == 0) + { + int acc = 0; + cu_kv[0] = 0; + for (int b = 0; b < batch; ++b) + { + acc += min(kv_lens[b], staged[b]); + cu_kv[b + 1] = acc; + } + } + for (int t = threadIdx.x; t < total_q; t += blockDim.x) + { + int const r = q_batch_row[t]; + int const pos = min(kv_lens[r], staged[r]) - qo[r] + q_intra[t]; + int const ps = max(pos, 0); + int const col = min(ps / page, ncols - 1); + out_loc[t] = table[r * table_ld + col] * page + ps % page; + if (n_valid != nullptr) + n_valid[t] = (max(pos + 1, 0) + page - 1) / page; + } +} +} // namespace + +cudaError_t invokeMinimaxM3MsaKvUpdate(int32_t const* kvLens, int32_t const* staged, int32_t const* qo, + int32_t const* qBatchRow, int32_t const* qIntra, int32_t const* table, int64_t tableLd, int32_t ncols, int32_t page, + int32_t batch, int32_t totalQ, int32_t* seqLensOut, int32_t* cuKv, int32_t* outLoc, int32_t* nValid, + cudaStream_t stream) +{ + msa_kv_update_kernel<<<1, 256, 0, stream>>>(kvLens, staged, qo, qBatchRow, qIntra, table, + static_cast(tableLd), ncols, page, batch, totalQ, seqLensOut, cuKv, outLoc, nValid); + return cudaGetLastError(); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/msaKvUpdate.h b/cpp/tensorrt_llm/kernels/minimaxM3/msaKvUpdate.h new file mode 100644 index 000000000000..d703cef6df73 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/msaKvUpdate.h @@ -0,0 +1,45 @@ +/* + * 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 +{ + +//! MiniMax-M3 MSA KV-length correction in one single-CTA launch, for all int32 operands: +//! +//! kvTrue[b] = min(kvLens[b], staged[b]) -> seqLensOut[b]; cuKv = exclusive prefix sum (optional) +//! pos[t] = kvTrue[r] - qo[r] + qIntra[t] (r = qBatchRow[t]) +//! outLoc[t] = table[r, min(max(pos, 0) / page, ncols - 1)] * page + max(pos, 0) % page +//! nValid[t] = (max(pos + 1, 0) + page - 1) / page (optional) +//! +//! cuKv and nValid may be null. table is row-major with row stride tableLd (elements) and ncols columns. +cudaError_t invokeMinimaxM3MsaKvUpdate(int32_t const* kvLens, int32_t const* staged, int32_t const* qo, + int32_t const* qBatchRow, int32_t const* qIntra, int32_t const* table, int64_t tableLd, int32_t ncols, int32_t page, + int32_t batch, int32_t totalQ, int32_t* seqLensOut, int32_t* cuKv, int32_t* outLoc, int32_t* nValid, + cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/mxfp8LowMGemm.cu b/cpp/tensorrt_llm/kernels/minimaxM3/mxfp8LowMGemm.cu new file mode 100644 index 000000000000..2e942d5b5a57 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/mxfp8LowMGemm.cu @@ -0,0 +1,820 @@ +/* + * 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/minimaxM3/mxfp8LowMGemm.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmException.h" + +#include +#include +#include +#include + +// The kernels use sm_100-family features (thread-block clusters, cp.async.bulk, st.async, mbarrier transaction +// counts, block-scaled conversions). Every other architecture of the build compiles them to empty bodies; the host +// entry points refuse to launch on those devices. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030) \ + && (defined(__CUDA_ARCH_FEAT_SM100_ALL) || defined(__CUDA_ARCH_FEAT_SM103_ALL) \ + || defined(__CUDA_ARCH_FAMILY_SPECIFIC__)) +#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 kMaxM = 48; +constexpr int kRowPad = 32; +constexpr int kEmptyWord = static_cast(0x80000000u); + +// Layer-boundary reduction inputs of the AR prologue (see m3_ar_fused.py); unused otherwise. +struct ArArgs +{ + int const* uc; + int const* cur_buf; + int64_t buf_words; + int world; + __nv_bfloat16 const* residual; + __nv_bfloat16 const* norm_w; + float eps; + __nv_bfloat16* residual_out; + uint8_t* xq_g; // [8][K] quantized activation published by tile 0 + uint8_t* xc_g; // [8][K / 32] its UE8M0 scales + int* flags; // [CLUSTER_K] per K-slice publish generation (call + 1) +}; + +// Optional L2 prefetch of the paged KV blocks the step's sparse decode will read (issued before the PDL wait, so it +// overlaps the producer kernel); all-null when unused. +struct PfArgs +{ + int const* block_table; // [rows, bt_stride] page ids + int const* seq_lens; // [rows] KV length per row + int rows; + int max_blocks; + int bt_stride; + char const* data; // page-major KV data pool + long long data_stride; // bytes between pages + int data_bytes; // contiguous bytes per page + char const* scale; // page-major block-scale pool (or null) + long long scale_stride; + int scale_bytes; +}; + +// Already-quantized activation (E4M3 rows [M, K] + UE8M0 scales in the 128x4-swizzled layout of +// trtllm::mxfp8_quantize(x, True)); all-null to quantize the BF16 activation in-kernel. +struct PreqArgs +{ + uint8_t const* xq; + uint8_t const* xsf; + int sf_k_tiles; // ceil(K / 128) +}; + +__device__ __forceinline__ void prefetch_kv_pages(PfArgs const& pf) +{ + if (pf.block_table == nullptr) + return; + int const units = pf.rows * pf.max_blocks; + for (int u = blockIdx.x; u < units; u += gridDim.x) + { + int const r = u / pf.max_blocks; + int const b = u - r * pf.max_blocks; + if (b * 128 >= pf.seq_lens[r]) + continue; + long long const page = pf.block_table[static_cast(r) * pf.bt_stride + b]; + char const* d = pf.data + page * pf.data_stride; + for (int off = threadIdx.x * 128; off < pf.data_bytes; off += blockDim.x * 128) + asm volatile("prefetch.global.L2::evict_last [%0];" ::"l"(d + off)); + if (pf.scale != nullptr) + { + char const* sc = pf.scale + page * pf.scale_stride; + for (int off = threadIdx.x * 128; off < pf.scale_bytes; off += blockDim.x * 128) + asm volatile("prefetch.global.L2::evict_last [%0];" ::"l"(sc + off)); + } + } +} + +__device__ __forceinline__ uint32_t mapa_shared(uint32_t addr, uint32_t target_rank) +{ + uint32_t r; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n" : "=r"(r) : "r"(addr), "r"(target_rank)); + return r; +} + +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__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 pow2_from_ue8m0(uint32_t e) +{ + return __uint_as_float(e == 0 ? 0x00400000u : (e << 23)); +} + +__device__ __forceinline__ void mma_e4m3( + 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.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\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), "f"(0.f), "f"(0.f), "f"(0.f), "f"(0.f)); +} + +__device__ __forceinline__ void bulk_g2s(void* dst, void const* src, uint32_t bytes, uint64_t* bar) +{ + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n" ::"r"( + smem_u32(dst)), + "l"(src), "r"(bytes), "r"(smem_u32(bar)) + : "memory"); +} + +__device__ __forceinline__ void mbar_wait(uint64_t* bar, uint32_t parity) +{ + asm volatile( + "{\n .reg .pred p;\n WAIT_%=:\n mbarrier.try_wait.parity.shared::cta.b64 p, [%0], %1;\n @!p bra " + "WAIT_%=;\n}\n" ::"r"(smem_u32(bar)), + "r"(parity) + : "memory"); +} + +// grid: (N / (16 * RT) * CLUSTER_K); cluster (CLUSTER_K, 1, 1); block WARPS warps. +// RT: 16-row weight tiles per CTA (amortizes the activation fetch + quantization over more rows). +// NT: number of 8-row activation tiles (MMA n=8 each), 1 for M <= 8 and 2 for M <= 16. +template +__global__ void __launch_bounds__(WARPS * 32) mxfp8_low_m_gemm_tma_kernel(__nv_bfloat16 const* __restrict__ x, int ldx, + int M, int N, int K, uint8_t const* __restrict__ w, uint8_t const* __restrict__ wsf, int num_k_tiles, + __nv_bfloat16* __restrict__ out, int ldo, int red_aliased, ArArgs ar, PfArgs pf, PreqArgs preq) +{ +#if MINIMAX_M3_KERNELS_ENABLED + constexpr int kThreads = WARPS * 32; + constexpr int kRows = 16 * RT; + constexpr int kTile = kRows * 8 * NT; + extern __shared__ __align__(128) uint8_t smem[]; + __shared__ __align__(8) uint64_t bar; + __shared__ __align__(8) uint64_t red_bar; + + int const rank = CLUSTER_K > 1 ? static_cast(cg::this_cluster().block_rank()) : 0; + int const tile = blockIdx.x / CLUSTER_K; + int const row0 = tile * kRows; + int const kc = K / CLUSTER_K; // columns owned by this CTA + int const kc_blocks = kc / 32; + int const k_begin = rank * kc; + int const kb_begin = k_begin / 32; + int const row_stride = kc + kRowPad; + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + + uint8_t* w_s = smem; // [kRows][row_stride] + uint8_t* sf_s = w_s + kRows * row_stride; // [RT][kc_blocks / 4][256] + uint8_t* xq_s = sf_s + RT * (kc_blocks / 4) * 256; // [M][kc] + uint8_t* xc_s = xq_s + M * kc; // [M][kc_blocks] + // Peer partial tiles (root only), then the warp partials unless they reuse the dead weight tile. + float* cluster_red = reinterpret_cast(xc_s + ((M * kc_blocks + 15) / 16) * 16); + float* red = red_aliased ? reinterpret_cast(w_s) : cluster_red + (CLUSTER_K - 1) * kTile; + // AR: the reduced + residual-added activation slice, BF16 [M][kc]. + __nv_bfloat16* u_s = reinterpret_cast<__nv_bfloat16*>( + reinterpret_cast(cluster_red + (CLUSTER_K - 1) * kTile) + (red_aliased ? 0 : WARPS * kTile * 4)); + __shared__ float tok_sq[kMaxM]; + __shared__ float peer_sq[CLUSTER_K][kMaxM]; + + uint32_t const w_bytes = kRows * kc; + uint32_t const sf_bytes = RT * (kc_blocks / 4) * 256u; + if (threadIdx.x == 0) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;\n" ::"r"(smem_u32(&bar))); + if (CLUSTER_K > 1 && rank == 0) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;\n" ::"r"(smem_u32(&red_bar))); + } + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + asm volatile( + "mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smem_u32(&bar)), "r"(w_bytes + sf_bytes) + : "memory"); + if (CLUSTER_K > 1 && rank == 0) + { + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smem_u32(&red_bar)), + "r"(static_cast((CLUSTER_K - 1) * kTile * sizeof(float))) + : "memory"); + } + } + __syncthreads(); + // The root's reduction barrier must be initialized before any peer pushes into it. + if constexpr (CLUSTER_K > 1) + { + asm volatile("barrier.cluster.arrive.release.aligned;\n" ::: "memory"); + } + // Weight rows and scale chunks do not depend on the producer kernel. + if (warp == 0) + { + for (int r = lane; r < kRows; r += 32) + { + bulk_g2s(w_s + r * row_stride, w + static_cast(row0 + r) * K + k_begin, kc, &bar); + } + int const chunks = kc_blocks / 4; + for (int t = lane; t < RT * chunks; t += 32) + { + int const rt = t / chunks; + int const c = t % chunks; + int const rr = row0 + 16 * rt; + int64_t const sf_base = static_cast(rr / 128) * num_k_tiles * 512 + (rr % 32) * 16; + bulk_g2s(sf_s + t * 256, wsf + sf_base + static_cast(kb_begin / 4 + c) * 512, 256, &bar); + } + } + prefetch_kv_pages(pf); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + if constexpr (AR) + { + // Layer-boundary reduction: sum the ranks' pushed partials for this CTA's K slice, add the residual + // (BF16 round, as the unfused add), then normalize with the RMS statistic of the whole row (the + // cluster spans K) and quantize like the unfused path. + if (threadIdx.x < kMaxM) + tok_sq[threadIdx.x] = 0.f; + __syncthreads(); + int const call = *reinterpret_cast(ar.cur_buf); + if (tile == 0) + { + int* base = const_cast(ar.uc) + static_cast(call & 1) * ar.buf_words; + int const half_k = K / 2; + for (int p = threadIdx.x; p < M * kc_blocks; p += kThreads) + { + int const m = p / kc_blocks; + int const b = p % kc_blocks; + int const k0 = k_begin + b * 32; + float v[32]; +#pragma unroll + for (int e = 0; e < 32; ++e) + v[e] = 0.f; + constexpr int kWorld = 4; + int4 wv[kWorld][4]; + int* src0 = base + static_cast(m) * kWorld * half_k + k0 / 2; +#pragma unroll + for (int r = 0; r < kWorld; ++r) + { +#pragma unroll + for (int q = 0; q < 4; ++q) + { + asm volatile("ld.volatile.global.v4.s32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(wv[r][q].x), "=r"(wv[r][q].y), "=r"(wv[r][q].z), "=r"(wv[r][q].w) + : "l"(src0 + r * half_k + 4 * q) + : "memory"); + } + } + long long spins = 0; + while (true) + { + bool pending = false; +#pragma unroll + for (int r = 0; r < kWorld; ++r) + { +#pragma unroll + for (int q = 0; q < 4; ++q) + { + int4 const c = wv[r][q]; + if (c.x == kEmptyWord || c.y == kEmptyWord || c.z == kEmptyWord || c.w == kEmptyWord) + { + pending = true; + asm volatile("ld.volatile.global.v4.s32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(wv[r][q].x), "=r"(wv[r][q].y), "=r"(wv[r][q].z), "=r"(wv[r][q].w) + : "l"(src0 + r * half_k + 4 * q) + : "memory"); + } + } + } + if (!pending) + break; + if (++spins > (1LL << 30)) + __trap(); + } +#pragma unroll + for (int r = 0; r < kWorld; ++r) + { +#pragma unroll + for (int q = 0; q < 4; ++q) + { + *reinterpret_cast(src0 + r * half_k + 4 * q) + = make_int4(kEmptyWord, kEmptyWord, kEmptyWord, kEmptyWord); + uint32_t const ww[4] = {static_cast(wv[r][q].x), static_cast(wv[r][q].y), + static_cast(wv[r][q].z), static_cast(wv[r][q].w)}; +#pragma unroll + for (int h = 0; h < 4; ++h) + { + v[q * 8 + 2 * h] += __uint_as_float(ww[h] << 16); + v[q * 8 + 2 * h + 1] += __uint_as_float(ww[h] & 0xFFFF0000u); + } + } + } + uint4 const* rsrc = reinterpret_cast(ar.residual + static_cast(m) * K + k0); + uint4* rdst = reinterpret_cast(ar.residual_out + static_cast(m) * K + k0); + uint4* udst = reinterpret_cast(u_s + m * kc + b * 32); + float sq = 0.f; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + uint4 const rraw = rsrc[q]; + __nv_bfloat162 const* rr = reinterpret_cast<__nv_bfloat162 const*>(&rraw); + uint32_t uw[4]; +#pragma unroll + for (int h = 0; h < 4; ++h) + { + float2 const rf = __bfloat1622float2(rr[h]); + __nv_bfloat162 const u + = __floats2bfloat162_rn(v[q * 8 + 2 * h] + rf.x, v[q * 8 + 2 * h + 1] + rf.y); + uw[h] = *reinterpret_cast(&u); + float2 const uf = __bfloat1622float2(u); + sq += uf.x * uf.x + uf.y * uf.y; + } + uint4 const uu = make_uint4(uw[0], uw[1], uw[2], uw[3]); + udst[q] = uu; + if (tile == 0) + rdst[q] = uu; + } + atomicAdd(&tok_sq[m], sq); + } + __syncthreads(); + if constexpr (CLUSTER_K > 1) + { + cg::cluster_group cl = cg::this_cluster(); + asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory"); + if (threadIdx.x < M) + { + for (int c = 0; c < CLUSTER_K; ++c) + { + *cl.map_shared_rank(&peer_sq[rank][threadIdx.x], c) = tok_sq[threadIdx.x]; + } + } + asm volatile("barrier.cluster.arrive.release.aligned;\n" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory"); + // Re-arm for the split-K reduction's wait below. + asm volatile("barrier.cluster.arrive.release.aligned;\n" ::: "memory"); + } + else + { + if (threadIdx.x < M) + peer_sq[0][threadIdx.x] = tok_sq[threadIdx.x]; + __syncthreads(); + } + for (int p = threadIdx.x; p < M * kc_blocks; p += kThreads) + { + int const m = p / kc_blocks; + int const b = p % kc_blocks; + float total = 0.f; +#pragma unroll + for (int c = 0; c < CLUSTER_K; ++c) + total += peer_sq[c][m]; + float const inv = rsqrtf(total / static_cast(K) + ar.eps); + uint4 const* usrc = reinterpret_cast(u_s + m * kc + b * 32); + uint4 const* wsrc = reinterpret_cast(ar.norm_w + k_begin + b * 32); + float v[32]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + uint4 const uraw = usrc[q]; + uint4 const wraw = wsrc[q]; + __nv_bfloat162 const* uu = reinterpret_cast<__nv_bfloat162 const*>(&uraw); + __nv_bfloat162 const* wwv = reinterpret_cast<__nv_bfloat162 const*>(&wraw); +#pragma unroll + for (int h = 0; h < 4; ++h) + { + float2 const uf = __bfloat1622float2(uu[h]); + float2 const wf = __bfloat1622float2(wwv[h]); + float2 const nf + = __bfloat1622float2(__floats2bfloat162_rn(uf.x * inv * wf.x, uf.y * inv * wf.y)); + v[q * 8 + 2 * h] = nf.x; + v[q * 8 + 2 * h + 1] = nf.y; + } + } + float amax = 0.f; +#pragma unroll + for (int t = 0; t < 32; ++t) + amax = fmaxf(amax, fabsf(v[t])); + float const sf = amax * rcp_approx_ftz(448.0f); + __nv_fp8_e8m0 code; + code.__x = __nv_cvt_float_to_e8m0(sf, __NV_SATFINITE, cudaRoundPosInf); + float const out_scale = amax != 0.f ? rcp_approx_ftz(static_cast(code)) : 0.0f; + uint32_t packed[8]; +#pragma unroll + for (int t = 0; t < 8; ++t) + { + __nv_fp8x2_e4m3 const lo(make_float2(v[4 * t] * out_scale, v[4 * t + 1] * out_scale)); + __nv_fp8x2_e4m3 const hi(make_float2(v[4 * t + 2] * out_scale, v[4 * t + 3] * out_scale)); + packed[t] = static_cast(lo.__x) | (static_cast(hi.__x) << 16); + } + uint4* dst = reinterpret_cast(xq_s + m * kc + b * 32); + dst[0] = make_uint4(packed[0], packed[1], packed[2], packed[3]); + dst[1] = make_uint4(packed[4], packed[5], packed[6], packed[7]); + xc_s[m * kc_blocks + b] = code.__x; + } + // Publish the quantized slice for the other tiles. + for (int p = threadIdx.x; p < M * kc_blocks; p += kThreads) + { + int const m = p / kc_blocks; + int const bb = p % kc_blocks; + uint4 const* src = reinterpret_cast(xq_s + m * kc + bb * 32); + uint4* dst = reinterpret_cast(ar.xq_g + static_cast(m) * K + k_begin + bb * 32); + dst[0] = src[0]; + dst[1] = src[1]; + ar.xc_g[static_cast(m) * (K / 32) + kb_begin + bb] = xc_s[m * kc_blocks + bb]; + } + __syncthreads(); + if (threadIdx.x == 0) + { + asm volatile("fence.acq_rel.gpu;" ::: "memory"); + asm volatile("st.release.gpu.global.s32 [%0], %1;" ::"l"(ar.flags + rank), "r"(call + 1) : "memory"); + } + } + else + { + if (threadIdx.x == 0) + { + int v; + long long spins = 0; + do + { + asm volatile("ld.acquire.gpu.global.s32 %0, [%1];" : "=r"(v) : "l"(ar.flags + rank) : "memory"); + if (++spins > (1LL << 31)) + __trap(); + } while (v - (call + 1) < 0); + } + __syncthreads(); + for (int p = threadIdx.x; p < M * kc_blocks; p += kThreads) + { + int const m = p / kc_blocks; + int const bb = p % kc_blocks; + uint4 const* src + = reinterpret_cast(ar.xq_g + static_cast(m) * K + k_begin + bb * 32); + uint4* dst = reinterpret_cast(xq_s + m * kc + bb * 32); + dst[0] = __ldcg(src); + dst[1] = __ldcg(src + 1); + xc_s[m * kc_blocks + bb] = __ldcg(ar.xc_g + static_cast(m) * (K / 32) + kb_begin + bb); + } + if constexpr (CLUSTER_K > 1) + { + // Keep the cluster barrier phases aligned with tile 0's prologue. + asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory"); + asm volatile("barrier.cluster.arrive.release.aligned;\n" ::: "memory"); + } + } + } + else if (preq.xq != nullptr) + { + // Stage this CTA's K slice of the pre-quantized activation. + int const vec_per_row = kc / 16; + int const vecs = M * vec_per_row; +#pragma unroll 4 + for (int p = threadIdx.x; p < vecs; p += kThreads) + { + int const m = p / vec_per_row; + int const v = p - m * vec_per_row; + reinterpret_cast(xq_s + m * kc)[v] + = __ldcg(reinterpret_cast(preq.xq + static_cast(m) * K + k_begin) + v); + } + for (int p = threadIdx.x; p < M * kc_blocks; p += kThreads) + { + int const m = p / kc_blocks; + int const b = p - m * kc_blocks; + int const kb = kb_begin + b; + xc_s[m * kc_blocks + b] = __ldcg(preq.xsf + static_cast(m / 128) * preq.sf_k_tiles * 512 + + (kb / 4) * 512 + (m % 32) * 16 + ((m % 128) / 32) * 4 + kb % 4); + } + } + else + { + // Quantize the activation columns owned by this CTA (native mxfp8_quantize math). Each thread + // issues the loads of all its blocks before converting any, so the CTA pays one L2 round trip. + int const items = M * kc_blocks; + for (int p0 = threadIdx.x; p0 < items; p0 += 2 * kThreads) + { + uint4 raw[2][4]; +#pragma unroll + for (int u = 0; u < 2; ++u) + { + int const p = p0 + u * kThreads; + if (p < items) + { + int const m = p / kc_blocks; + int const b = p % kc_blocks; + uint4 const* src + = reinterpret_cast(x + static_cast(m) * ldx + k_begin + b * 32); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + raw[u][q] = src[q]; + } + } + } +#pragma unroll + for (int u = 0; u < 2; ++u) + { + int const p = p0 + u * kThreads; + if (p >= items) + { + continue; + } + int const m = p / kc_blocks; + int const b = p % kc_blocks; + float v[32]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&raw[u][q]); +#pragma unroll + for (int t = 0; t < 4; ++t) + { + float2 const f = __bfloat1622float2(h[t]); + v[q * 8 + t * 2] = f.x; + v[q * 8 + t * 2 + 1] = f.y; + } + } + float amax = 0.f; +#pragma unroll + for (int t = 0; t < 32; ++t) + { + amax = fmaxf(amax, fabsf(v[t])); + } + float const sf = amax * rcp_approx_ftz(448.0f); + __nv_fp8_e8m0 code; + code.__x = __nv_cvt_float_to_e8m0(sf, __NV_SATFINITE, cudaRoundPosInf); + float const out_scale = amax != 0.f ? rcp_approx_ftz(static_cast(code)) : 0.0f; + uint32_t packed[8]; +#pragma unroll + for (int t = 0; t < 8; ++t) + { + __nv_fp8x2_e4m3 const lo(make_float2(v[4 * t] * out_scale, v[4 * t + 1] * out_scale)); + __nv_fp8x2_e4m3 const hi(make_float2(v[4 * t + 2] * out_scale, v[4 * t + 3] * out_scale)); + packed[t] = static_cast(lo.__x) | (static_cast(hi.__x) << 16); + } + uint4* dst = reinterpret_cast(xq_s + m * kc + b * 32); + dst[0] = make_uint4(packed[0], packed[1], packed[2], packed[3]); + dst[1] = make_uint4(packed[4], packed[5], packed[6], packed[7]); + xc_s[m * kc_blocks + b] = code.__x; + } + } + } + __syncthreads(); + mbar_wait(&bar, 0); + + int const g = lane >> 2; + int const j = lane & 3; + int const col0 = 2 * j; + int const bpw = kc_blocks / WARPS; + int const chunks = kc_blocks / 4; + float acc[RT][NT][4] = {}; +#pragma unroll 2 + for (int i = 0; i < bpw; ++i) + { + int const b = warp * bpw + i; + int const kb = kb_begin + b; + uint2 xb[NT]; + float xs0[NT], xs1[NT]; +#pragma unroll + for (int t = 0; t < NT; ++t) + { + int const row = 8 * t + g; + int const c0 = 8 * t + col0; + xb[t] = row < M ? *reinterpret_cast(xq_s + row * kc + b * 32 + 8 * j) : make_uint2(0, 0); + xs0[t] = c0 < M ? pow2_from_ue8m0(xc_s[c0 * kc_blocks + b]) : 0.f; + xs1[t] = c0 + 1 < M ? pow2_from_ue8m0(xc_s[(c0 + 1) * kc_blocks + b]) : 0.f; + } +#pragma unroll + for (int rt = 0; rt < RT; ++rt) + { + uint8_t const* wa_s = w_s + (16 * rt + g) * row_stride + 8 * j + b * 32; + uint2 const wa = *reinterpret_cast(wa_s); + uint2 const wb = *reinterpret_cast(wa_s + 8 * row_stride); + int const inner_m = ((row0 + 16 * rt) % 128) / 32; + uint8_t const* sft = sf_s + (rt * chunks + b / 4) * 256 + inner_m * 4 + kb % 4; + float const wsa = pow2_from_ue8m0(sft[g * 16]); + float const wsb = pow2_from_ue8m0(sft[(g + 8) * 16]); +#pragma unroll + for (int t = 0; t < NT; ++t) + { + float d[4]; + mma_e4m3(d, wa.x, wb.x, wa.y, wb.y, xb[t].x, xb[t].y); + acc[rt][t][0] += d[0] * (wsa * xs0[t]); + acc[rt][t][1] += d[1] * (wsa * xs1[t]); + acc[rt][t][2] += d[2] * (wsb * xs0[t]); + acc[rt][t][3] += d[3] * (wsb * xs1[t]); + } + } + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + if (red_aliased) + { + __syncthreads(); + } +#pragma unroll + for (int rt = 0; rt < RT; ++rt) + { +#pragma unroll + for (int t = 0; t < NT; ++t) + { + int const ra = 16 * rt + g; + red[warp * kTile + ra * 8 * NT + 8 * t + col0] = acc[rt][t][0]; + red[warp * kTile + ra * 8 * NT + 8 * t + col0 + 1] = acc[rt][t][1]; + red[warp * kTile + (ra + 8) * 8 * NT + 8 * t + col0] = acc[rt][t][2]; + red[warp * kTile + (ra + 8) * 8 * NT + 8 * t + col0 + 1] = acc[rt][t][3]; + } + } + __syncthreads(); + + // Warp reduction in fixed order. Peers push their partial tile into the root's shared memory with + // st.async (completing bytes on the root's mbarrier) and exit; the root sums ranks in order. + if constexpr (CLUSTER_K > 1) + { + asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory"); + } + if (CLUSTER_K == 1 || rank == 0) + { + if constexpr (CLUSTER_K > 1) + { + uint32_t ready = 0; + while (!ready) + { + 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}\n" + : "=r"(ready) + : "r"(smem_u32(&red_bar)) + : "memory"); + } + } + __syncthreads(); + for (int idx = threadIdx.x; idx < kTile; idx += kThreads) + { + int const r = idx / (8 * NT); + int const m = idx % (8 * NT); + if (m < M) + { + float total = 0.f; +#pragma unroll + for (int ww = 0; ww < WARPS; ++ww) + { + total += red[ww * kTile + idx]; + } +#pragma unroll + for (int c = 1; c < CLUSTER_K; ++c) + { + total += cluster_red[(c - 1) * kTile + idx]; + } + out[static_cast(m) * ldo + row0 + r] = __float2bfloat16_rn(total); + } + } + } + else + { + uint32_t const remote_bar = mapa_shared(smem_u32(&red_bar), 0); + uint32_t const remote_dst = mapa_shared(smem_u32(cluster_red + (rank - 1) * kTile), 0); + for (int idx = threadIdx.x; idx < kTile; idx += kThreads) + { + float s = 0.f; +#pragma unroll + for (int ww = 0; ww < WARPS; ++ww) + { + s += red[ww * kTile + idx]; + } + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + remote_dst + idx * 4), + "r"(__float_as_uint(s)), "r"(remote_bar) + : "memory"); + } + } +#endif +} + +template +void launch(__nv_bfloat16 const* x, int ldx, int M, int N, int K, uint8_t const* w, uint8_t const* wsf, + __nv_bfloat16* out, int ldo, bool pdl, cudaStream_t stream) +{ + constexpr bool AR = false; + int const kc = K / CLUSTER_K; + int const kc_blocks = kc / 32; + TLLM_CHECK_WITH_INFO(N % (16 * RT) == 0 && K % (CLUSTER_K * 128) == 0 && kc_blocks % WARPS == 0, + "unsupported shape N=%d K=%d", N, K); + constexpr int kTile = 16 * RT * 8 * NT; + size_t const w_bytes = 16 * RT * (kc + kRowPad); + size_t const red_bytes = static_cast(WARPS) * kTile * sizeof(float); + int const red_aliased = red_bytes <= w_bytes ? 1 : 0; + size_t const smem = w_bytes + RT * (kc_blocks / 4) * 256 + M * kc + ((M * kc_blocks + 15) / 16) * 16 + + (CLUSTER_K - 1) * kTile * sizeof(float) + (red_aliased ? 0 : red_bytes); + auto kernel = mxfp8_low_m_gemm_tma_kernel; + // Configure once, outside any CUDA graph capture, to the largest supported request. + static bool const configured = [&]() + { + int device = 0, max_optin = 0; + cudaGetDevice(&device); + cudaDeviceGetAttribute(&max_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, device); + cudaFuncAttributes attrs = {}; + if (cudaFuncGetAttributes(&attrs, kernel) != cudaSuccess) + { + cudaGetLastError(); + return false; + } + bool const ok = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + max_optin - static_cast(attrs.sharedSizeBytes)) + == cudaSuccess; + if (!ok) + { + cudaGetLastError(); + } + return ok; + }(); + TLLM_CHECK_WITH_INFO(configured, "failed to configure dynamic shared memory"); + cudaFuncAttributes fattrs = {}; + cudaFuncGetAttributes(&fattrs, kernel); + TLLM_CHECK_WITH_INFO(static_cast(smem) <= fattrs.maxDynamicSharedSizeBytes, + "shared memory request exceeds the configured limit"); + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(N / (16 * RT) * CLUSTER_K); + cfg.blockDim = dim3(WARPS * 32); + 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 = CLUSTER_K; + 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, x, ldx, M, N, K, w, wsf, (K / 32 + 3) / 4, out, ldo, red_aliased, ArArgs{}, PfArgs{}, PreqArgs{}); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +} // namespace + +bool isMinimaxM3Mxfp8LowMGemmSupported(int32_t warps, int32_t clusterK, int32_t rowTiles, int32_t m) +{ + bool const built = (warps == 4 && clusterK == 4 && rowTiles == 2) || (warps == 8 && clusterK == 2 && rowTiles == 4) + || (warps == 4 && clusterK == 4 && rowTiles == 1) || (warps == 4 && clusterK == 2 && rowTiles == 2); + return built && m >= 1 && m <= 8; +} + +void invokeMinimaxM3Mxfp8LowMGemm(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t n, int32_t k, + uint8_t const* w, uint8_t const* wsf, __nv_bfloat16* out, int32_t ldo, int32_t warps, int32_t clusterK, + int32_t rowTiles, bool pdl, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(isMinimaxM3Mxfp8LowMGemmSupported(warps, clusterK, rowTiles, m), + "unsupported configuration (warps %d, cluster_k %d, row_tiles %d) at M=%d", warps, clusterK, rowTiles, m); + // One 8-row activation tile (NT = 1) covers every built row count. + if (warps == 4 && clusterK == 4 && rowTiles == 2) + { + launch<4, 4, 1, 2>(x, ldx, m, n, k, w, wsf, out, ldo, pdl, stream); + } + else if (warps == 8 && clusterK == 2 && rowTiles == 4) + { + launch<8, 2, 1, 4>(x, ldx, m, n, k, w, wsf, out, ldo, pdl, stream); + } + else if (warps == 4 && clusterK == 4 && rowTiles == 1) + { + launch<4, 4, 1, 1>(x, ldx, m, n, k, w, wsf, out, ldo, pdl, stream); + } + else + { + launch<4, 2, 1, 2>(x, ldx, m, n, k, w, wsf, out, ldo, pdl, stream); + } +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/mxfp8LowMGemm.h b/cpp/tensorrt_llm/kernels/minimaxM3/mxfp8LowMGemm.h new file mode 100644 index 000000000000..433cb6a1fe07 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/mxfp8LowMGemm.h @@ -0,0 +1,46 @@ +/* + * 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 +{ + +//! Whether `invokeMinimaxM3Mxfp8LowMGemm` is built for this (warps, clusterK, rowTiles) configuration at M rows: +//! (4, 4, 2), (8, 2, 4), (4, 4, 1) and (4, 2, 2), each for 1 <= M <= 8. +bool isMinimaxM3Mxfp8LowMGemmSupported(int32_t warps, int32_t clusterK, int32_t rowTiles, int32_t m); + +//! Fused MXFP8 activation quantization + low-M MXFP8 GEMM: `out[M, N] = mxfp8(x)[M, K] @ w[N, K]^T` in BF16. +//! +//! `x` is BF16 [M, K] with row stride `ldx`; `w` is E4M3 [N, K] row-major with UE8M0 scales `wsf` in the 128x4 +//! swizzled layout of `mxfp8_quantize(w, True)`; `out` is BF16 [M, N] with row stride `ldo`. Each CTA owns +//! 16 * rowTiles weight rows and K / clusterK columns; `warps` warps per CTA. +void invokeMinimaxM3Mxfp8LowMGemm(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t n, int32_t k, + uint8_t const* w, uint8_t const* wsf, __nv_bfloat16* out, int32_t ldo, int32_t warps, int32_t clusterK, + int32_t rowTiles, bool pdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/producerPw.cu b/cpp/tensorrt_llm/kernels/minimaxM3/producerPw.cu new file mode 100644 index 000000000000..fb3e4bfb3a89 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/producerPw.cu @@ -0,0 +1,324 @@ +/* + * 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/minimaxM3/producerPw.h" + +#include +#include +#include + +// MiniMax-M3 NVFP4 QKV + indexer norm/RoPE/KV-insert producer with an early PDL trigger. +// +// Same math and outputs as minimaxM3Nvfp4QKVIndexerNormRopeKVInsertKernel (fusedQKNormRopeKernel.cu, built with +// --use_fast_math like the rest of the library), specialized to 1 KV head, 1 index head, head_dim 128, rotary_dim 64, +// page 128, NVFP4 main KV and FP8 index-K, with 8 (TP8) or 16 (TP4) query heads per rank. One warp per (token, +// packed head). +// +// 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 `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 minimaxM3 +{ +namespace producerPw +{ +namespace +{ + +constexpr int kHeadDim = 128; +constexpr int kRotaryDim = 64; +constexpr int kPageSize = 128; +constexpr int kElems = kHeadDim / 32; +constexpr int kHeadsKV = 1; +constexpr int kHeadsIndex = 1; +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; + + // Only constant data is 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; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + // Position ids and cache slots are per-step metadata, read after the wait. The packed row, the position id and + // the slot are independent loads, and the RoPE row load is issued before the norm reduction, so the + // native kernel's serial chain of dependent loads becomes two round trips. + uint2 const packed = *reinterpret_cast(p.packed + inputOffset); + bool const rotates = !isV && laneId * kElems < kRotaryDim; + int const positionId = rotates ? p.positionIds[tokenIdx] : 0; + int const slot = (isK || isV || isIndexK) ? p.outCacheLoc[tokenIdx] : 0; + float cosine[kElems] = {}; + float sine[kElems] = {}; + if (rotates) + { + int64_t const ropeRow = static_cast(positionId) * kRotaryDim; + int const coefficient = (laneId * kElems) % (kRotaryDim / 2); + float4 const c = *reinterpret_cast(p.rotaryCosSin + ropeRow + coefficient); + float4 const s = *reinterpret_cast(p.rotaryCosSin + ropeRow + kRotaryDim / 2 + coefficient); + cosine[0] = c.x; + cosine[1] = c.y; + cosine[2] = c.z; + cosine[3] = c.w; + sine[0] = s.x; + sine[1] = s.y; + sine[2] = s.z; + sine[3] = 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 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 + 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(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 invokeProducerPw(Params const& p, int numTokens, int headsQ, bool pdl, bool early, cudaStream_t stream) +{ + if (numTokens <= 0) + return cudaSuccess; + if (headsQ == 16) + return early ? launch<16, true>(p, numTokens, pdl, stream) : launch<16, false>(p, numTokens, pdl, stream); + if (headsQ == 8) + return early ? launch<8, true>(p, numTokens, pdl, stream) : launch<8, false>(p, numTokens, pdl, stream); + return cudaErrorInvalidValue; +} + +} // namespace producerPw +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/producerPw.h b/cpp/tensorrt_llm/kernels/minimaxM3/producerPw.h new file mode 100644 index 000000000000..6d6848e0e199 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/producerPw.h @@ -0,0 +1,69 @@ +/* + * 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 minimaxM3 +{ +//! The MiniMax-M3 NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer with its constant loads before the PDL +//! grid-dependency wait and an optional early dependents trigger (sm_103a). Same outputs, bit for bit, as +//! minimaxM3Nvfp4QKVIndexerNormRopeKVInsert. See producerPw.cu. +namespace producerPw +{ + +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; +}; + +//! Launch over `numTokens` tokens, one warp per (token, packed head); headsQ is 8 or 16. +cudaError_t invokeProducerPw(Params const& p, int numTokens, int headsQ, bool pdl, bool early, cudaStream_t stream); + +} // namespace producerPw +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.cu b/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.cu new file mode 100644 index 000000000000..d12b765e5ecf --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.cu @@ -0,0 +1,606 @@ +/* + * 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/minimaxM3/routerGemv.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" + +#include +#include +#include + +// The kernels use sm_100-family features (thread-block clusters, cp.async.bulk, st.async, mbarrier transaction +// counts, block-scaled conversions). Every other architecture of the build compiles them to empty bodies; the host +// entry points refuse to launch on those devices. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030) \ + && (defined(__CUDA_ARCH_FEAT_SM100_ALL) || defined(__CUDA_ARCH_FEAT_SM103_ALL) \ + || defined(__CUDA_ARCH_FAMILY_SPECIFIC__)) +#define MINIMAX_M3_KERNELS_ENABLED 1 +#else +#define MINIMAX_M3_KERNELS_ENABLED 0 +#endif + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr int kMaxM = 16; + +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void bulk_g2s(void* dst, void const* src, uint32_t bytes, uint64_t* bar) +{ + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n" ::"r"( + smem_u32(dst)), + "l"(src), "r"(bytes), "r"(smem_u32(bar)) + : "memory"); +} + +__device__ __forceinline__ void mbar_wait(uint64_t* bar, uint32_t parity) +{ + asm volatile( + "{\n .reg .pred p;\n WAIT_%=:\n mbarrier.try_wait.parity.shared::cta.b64 p, [%0], %1;\n @!p bra " + "WAIT_%=;\n}\n" ::"r"(smem_u32(bar)), + "r"(parity) + : "memory"); +} + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +__device__ __forceinline__ uint32_t e2m1x8(float2 const (&v)[4]) +{ + uint32_t r; + asm("{\n .reg .b8 b0, b1, b2, b3;\n" + " cvt.rn.satfinite.e2m1x2.f32 b0, %2, %1;\n cvt.rn.satfinite.e2m1x2.f32 b1, %4, %3;\n" + " cvt.rn.satfinite.e2m1x2.f32 b2, %6, %5;\n cvt.rn.satfinite.e2m1x2.f32 b3, %8, %7;\n" + " mov.b32 %0, {b0, b1, b2, b3};\n}\n" + : "=r"(r) + : "f"(v[0].x), "f"(v[0].y), "f"(v[1].x), "f"(v[1].y), "f"(v[2].x), "f"(v[2].y), "f"(v[3].x), "f"(v[3].y)); + return r; +} + +// NVFP4 (16-element E4M3 block scale, per-tensor global scale) quantization of one 16-value block, bit-identical +// to trtllm fp4_quantize (cvt_warp_fp16_to_fp4, linear scale layout). +__device__ __forceinline__ void quant_block16(__nv_bfloat16 const* src, float sf_scale, uint2* q, uint8_t* sf) +{ + uint4 const v0 = *reinterpret_cast(src); + uint4 const v1 = *reinterpret_cast(src + 8); + __nv_bfloat162 h[8]; + *reinterpret_cast(&h[0]) = v0; + *reinterpret_cast(&h[4]) = v1; + __nv_bfloat162 mx = __habs2(h[0]); +#pragma unroll + for (int i = 1; i < 8; ++i) + mx = __hmax2(mx, __habs2(h[i])); + float const vec_max = fmaxf(__low2float(mx), __high2float(mx)); + float sf_value = sf_scale * (vec_max * rcp_approx_ftz(6.0f)); + __nv_fp8_e4m3 const tmp = __nv_fp8_e4m3(sf_value); + sf_value = static_cast(tmp); + float const out_scale = vec_max != 0.f ? rcp_approx_ftz(sf_value * rcp_approx_ftz(sf_scale)) : 0.f; + float2 lo[4], hi[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + lo[i] = __bfloat1622float2(h[i]); + lo[i].x *= out_scale; + lo[i].y *= out_scale; + hi[i] = __bfloat1622float2(h[i + 4]); + hi[i].x *= out_scale; + hi[i].y *= out_scale; + } + *q = make_uint2(e2m1x8(lo), e2m1x8(hi)); + *sf = tmp.__x; +} + +// grid: E / ROWS; block: 256 threads. smem: ROWS * K fp32 weights + M * K bf16 activations. +// With q_out, the CTAs also emit the NVFP4 activation ([M, K / 2] packed, [M, K / 16] E4M3 scales) for the MoE. +template +__global__ void __launch_bounds__(256) m3_router_gemv_kernel(__nv_bfloat16 const* __restrict__ x, int ldx, + float const* __restrict__ w, float* __restrict__ out, int M, int K, int E, uint8_t* __restrict__ q_out, + uint8_t* __restrict__ sf_out, float const* __restrict__ sf_scale) +{ +#if MINIMAX_M3_KERNELS_ENABLED + extern __shared__ __align__(128) uint8_t smem[]; + __shared__ __align__(8) uint64_t w_bar; + __shared__ __align__(8) uint64_t x_bar; + __shared__ float red[8][ROWS][kMaxM]; + float* w_s = reinterpret_cast(smem); + __nv_bfloat16* x_s = reinterpret_cast<__nv_bfloat16*>(smem + ROWS * K * sizeof(float)); + int const e0 = blockIdx.x * ROWS; + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + if (threadIdx.x == 0) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;\n" ::"r"(smem_u32(&w_bar))); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;\n" ::"r"(smem_u32(&x_bar))); + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smem_u32(&w_bar)), + "r"(static_cast(ROWS * K * sizeof(float))) + : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smem_u32(&x_bar)), + "r"(static_cast(M * K * sizeof(__nv_bfloat16))) + : "memory"); + } + __syncthreads(); + // Router weights do not depend on the producer kernel. + if (warp == 0 && lane < ROWS) + { + bulk_g2s(w_s + lane * K, w + static_cast(e0 + lane) * K, K * sizeof(float), &w_bar); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + if (warp == 1 && lane < M) + { + bulk_g2s(x_s + lane * K, x + static_cast(lane) * ldx, K * sizeof(__nv_bfloat16), &x_bar); + } + mbar_wait(&x_bar, 0); + if (q_out != nullptr) + { + float const s = *sf_scale; + int const kblocks = K / 16; + for (int b = blockIdx.x * blockDim.x + threadIdx.x; b < M * kblocks; b += gridDim.x * blockDim.x) + { + int const m = b / kblocks; + int const kb = b - m * kblocks; + quant_block16(x_s + m * K + kb * 16, s, + reinterpret_cast(q_out + static_cast(m) * (K / 2) + kb * 8), + sf_out + static_cast(m) * kblocks + kb); + } + } + mbar_wait(&w_bar, 0); + + float acc[ROWS][kMaxM]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) + { +#pragma unroll + for (int m = 0; m < kMaxM; ++m) + { + acc[r][m] = 0.f; + } + } + // Each thread owns 4 consecutive columns per 1024-column stride. + for (int k = threadIdx.x * 4; k < K; k += 256 * 4) + { + float4 wv[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) + { + wv[r] = *reinterpret_cast(w_s + r * K + k); + } +#pragma unroll + for (int m = 0; m < kMaxM; ++m) + { + if (m < M) + { + uint2 const raw = *reinterpret_cast(x_s + m * K + k); + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&raw); + float2 const a = __bfloat1622float2(h[0]); + float2 const b = __bfloat1622float2(h[1]); +#pragma unroll + for (int r = 0; r < ROWS; ++r) + { + acc[r][m] = fmaf(a.x, wv[r].x, acc[r][m]); + acc[r][m] = fmaf(a.y, wv[r].y, acc[r][m]); + acc[r][m] = fmaf(b.x, wv[r].z, acc[r][m]); + acc[r][m] = fmaf(b.y, wv[r].w, acc[r][m]); + } + } + } + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +#pragma unroll + for (int r = 0; r < ROWS; ++r) + { +#pragma unroll + for (int m = 0; m < kMaxM; ++m) + { + if (m < M) + { + float v = acc[r][m]; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + v += __shfl_xor_sync(0xFFFFFFFFu, v, o); + } + if (lane == 0) + { + red[warp][r][m] = v; + } + } + } + } + __syncthreads(); + if (threadIdx.x < ROWS * M) + { + int const r = threadIdx.x / M; + int const m = threadIdx.x % M; + float s = 0.f; +#pragma unroll + for (int ww = 0; ww < 8; ++ww) + { + s += red[ww][r][m]; + } + out[static_cast(m) * E + e0 + r] = s; + } +#endif +} + +// ---------------- 9..48 tokens: TF32 tensor-core router (same precision as the native TF32 GEMM) ---------------- +// grid: 8 expert tiles (16 experts) x CL CTAs, one CL-CTA cluster per tile, 256 threads. CTA r of a cluster computes +// the [16 experts x M] partials of its 16 / CL K slices (16 slices of 384), each over the same 8-warp split. CTA r +// owns 16 / CL experts of the tile: every CTA st.asyncs its partials for those experts into the owner's mailbox, and +// the owner sums the 16 slices in slice order, so the logits do not depend on CL. Optionally each CTA also emits the +// NVFP4 activation for its K range and the token rows m with m % 8 == expert tile. +constexpr int kTcExperts = 16; +constexpr int kTcSlices = 16; +constexpr int kTcMaxM = 48; + +__device__ __forceinline__ uint32_t f2tf32(float f) +{ + uint32_t r; + asm("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(r) : "f"(f)); + return r; +} + +__device__ __forceinline__ uint32_t mapa_u32(void const* p, uint32_t rank) +{ + uint32_t r; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n" : "=r"(r) : "r"(smem_u32(p)), "r"(rank)); + return r; +} + +template +__global__ void __launch_bounds__(256) m3_router_tc_kernel(__nv_bfloat16 const* __restrict__ x, int ldx, + float const* __restrict__ w, float* __restrict__ out, int M, int K, uint8_t* __restrict__ q_out, + uint8_t* __restrict__ sf_out, float const* __restrict__ sf_scale) +{ +#if MINIMAX_M3_KERNELS_ENABLED + extern __shared__ __align__(128) uint8_t smem[]; + __shared__ __align__(8) uint64_t w_bar; + __shared__ __align__(8) uint64_t x_bar; + __shared__ __align__(8) uint64_t mb_bar; + constexpr int SPC = kTcSlices / CL; // K slices per CTA + constexpr int EPC = kTcExperts / CL; // experts owned per CTA + __shared__ float red[8][kTcExperts][kTcMaxM]; + __shared__ float part[SPC][kTcExperts][kTcMaxM]; // this CTA's slice partials + __shared__ float mailbox[CL][SPC][EPC][kTcMaxM]; // [source CTA][slice][owned expert][m] + int const tile = blockIdx.x / CL; + int const crank = blockIdx.x % CL; + int const ks = K / kTcSlices; // 384 + int const kc = SPC * ks; + int const k0 = crank * kc; + int const e0 = tile * kTcExperts; + int const wst = kc + 4; // padded row strides (bank-conflict-free fragment loads) + int const xst = kc + 8; + float* w_s = reinterpret_cast(smem); // [16][wst] + __nv_bfloat16* x_s = reinterpret_cast<__nv_bfloat16*>(smem + kTcExperts * wst * 4); // [M][xst] + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + if (threadIdx.x == 0) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;\n" ::"r"(smem_u32(&w_bar))); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;\n" ::"r"(smem_u32(&x_bar))); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;\n" ::"r"(smem_u32(&mb_bar))); + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smem_u32(&w_bar)), + "r"(static_cast(kTcExperts * kc * 4)) + : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smem_u32(&x_bar)), + "r"(static_cast(M * kc * 2)) + : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smem_u32(&mb_bar)), + "r"(static_cast((CL - 1) * SPC * EPC * M * 4)) + : "memory"); + } + __syncthreads(); + // The peers' mailbox barriers must be initialized before this CTA's st.async lands there. + asm volatile("barrier.cluster.arrive.release.aligned;\n" ::: "memory"); + if (warp == 0 && lane < kTcExperts) + bulk_g2s(w_s + lane * wst, w + static_cast(e0 + lane) * K + k0, kc * 4, &w_bar); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + if (warp == 1) + { + for (int m = lane; m < M; m += 32) + bulk_g2s(x_s + m * xst, x + static_cast(m) * ldx + k0, kc * 2, &x_bar); + } + mbar_wait(&x_bar, 0); + if (q_out != nullptr) + { + float const sc = *sf_scale; + int const kb_cta = kc / 16; + for (int b = threadIdx.x; b < M * kb_cta; b += blockDim.x) + { + int const m = b / kb_cta; + if (m % 8 != tile) + continue; // 8 expert tiles share the row quantization + int const kb = b - m * kb_cta; + int const kbg = k0 / 16 + kb; + quant_block16(x_s + m * xst + kb * 16, sc, + reinterpret_cast(q_out + static_cast(m) * (K / 2) + kbg * 8), + sf_out + static_cast(m) * (K / 16) + kbg); + } + } + mbar_wait(&w_bar, 0); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + // Per slice h: warp w owns k8 steps [w * steps, (w + 1) * steps) of the slice; all token tiles (n = 8 each). + int const g = lane >> 2; + int const t = lane & 3; + int const ntiles = (M + 7) / 8; + int const steps = ks / 8 / 8; +#pragma unroll 1 + for (int h = 0; h < SPC; ++h) + { + float acc[6][4]; +#pragma unroll + for (int i = 0; i < 6; ++i) + acc[i][0] = acc[i][1] = acc[i][2] = acc[i][3] = 0.f; + for (int st = warp * steps; st < (warp + 1) * steps; ++st) + { + int const kk = h * ks + st * 8; + uint32_t const a0 = f2tf32(w_s[g * wst + kk + t]); + uint32_t const a1 = f2tf32(w_s[(g + 8) * wst + kk + t]); + uint32_t const a2 = f2tf32(w_s[g * wst + kk + t + 4]); + uint32_t const a3 = f2tf32(w_s[(g + 8) * wst + kk + t + 4]); +#pragma unroll + for (int nt = 0; nt < 6; ++nt) + { + if (nt < ntiles) + { + int const m = nt * 8 + g; + float bx0 = 0.f, bx1 = 0.f; + if (m < M) + { + bx0 = __bfloat162float(x_s[m * xst + kk + t]); + bx1 = __bfloat162float(x_s[m * xst + kk + t + 4]); + } + uint32_t const b0 = f2tf32(bx0), b1 = f2tf32(bx1); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9}, {%0,%1,%2,%3};\n" + : "+f"(acc[nt][0]), "+f"(acc[nt][1]), "+f"(acc[nt][2]), "+f"(acc[nt][3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); + } + } + } +#pragma unroll + for (int nt = 0; nt < 6; ++nt) + { + if (nt < ntiles) + { + int const c = nt * 8 + 2 * t; + red[warp][g][c] = acc[nt][0]; + red[warp][g][c + 1] = acc[nt][1]; + red[warp][g + 8][c] = acc[nt][2]; + red[warp][g + 8][c + 1] = acc[nt][3]; + } + } + __syncthreads(); + for (int i = threadIdx.x; i < kTcExperts * M; i += blockDim.x) + { + int const e = i / M, m = i % M; + float v = 0.f; +#pragma unroll + for (int ww = 0; ww < 8; ++ww) + v += red[ww][e][m]; + part[h][e][m] = v; + } + __syncthreads(); + } + asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory"); + // Expert e of the tile belongs to CTA e / EPC: its slice partials go to that CTA's mailbox[crank]. + for (int i = threadIdx.x; i < SPC * kTcExperts * M; i += blockDim.x) + { + int const h = i / (kTcExperts * M); + int const r = i - h * kTcExperts * M; + int const e = r / M, m = r % M; + int const owner = e / EPC; + float const v = part[h][e][m]; + float* dst = &mailbox[crank][h][e % EPC][m]; + if (owner == crank) + { + *dst = v; + } + else + { + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];\n" ::"r"( + mapa_u32(dst, owner)), + "r"(__float_as_uint(v)), "r"(mapa_u32(&mb_bar, owner)) + : "memory"); + } + } + __syncthreads(); + { + uint32_t done = 0; + while (!done) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}\n" + : "=r"(done) + : "r"(smem_u32(&mb_bar)) + : "memory"); + } + } + for (int i = threadIdx.x; i < EPC * M; i += blockDim.x) + { + int const le = i / M, m = i % M; + float v = 0.f; +#pragma unroll + for (int src = 0; src < CL; ++src) + { +#pragma unroll + for (int h = 0; h < SPC; ++h) + v += mailbox[src][h][le][m]; + } + out[static_cast(m) * (kTcExperts * 8) + e0 + EPC * crank + le] = v; + } +#endif +} + +template +void launchGemv(__nv_bfloat16 const* x, int ldx, int M, int K, float const* w, int E, float* out, uint8_t* q_out, + uint8_t* sf_out, float const* sf_scale, bool pdl, cudaStream_t stream) +{ + size_t const smem = static_cast(ROWS) * K * sizeof(float) + static_cast(M) * K * 2; + auto kernel = m3_router_gemv_kernel; + static bool const configured = [&]() + { + int device = 0, max_optin = 0; + cudaGetDevice(&device); + cudaDeviceGetAttribute(&max_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, device); + cudaFuncAttributes attrs = {}; + if (cudaFuncGetAttributes(&attrs, kernel) != cudaSuccess) + { + return false; + } + return cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + max_optin - static_cast(attrs.sharedSizeBytes)) + == cudaSuccess; + }(); + TLLM_CHECK_WITH_INFO(configured, "failed to configure dynamic shared memory"); + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(E / ROWS); + cfg.blockDim = dim3(256); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + cudaError_t const err = cudaLaunchKernelEx(&cfg, kernel, x, ldx, w, out, M, K, E, q_out, sf_out, sf_scale); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +} // namespace + +void invokeMinimaxM3RouterGemv(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t k, float const* w, int32_t e, + float* out, uint8_t* qOut, uint8_t* sfOut, float const* sfScale, bool pdl, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(m >= 1 && m <= kMaxM && k % 1024 == 0, "unsupported router GEMV shape M=%d K=%d", m, k); + 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) +{ + 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 = [&]() + { + int device = 0, max_optin = 0; + cudaGetDevice(&device); + cudaDeviceGetAttribute(&max_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, device); + auto configure = [&](auto kernel, bool non_portable) + { + cudaFuncAttributes attrs = {}; + if (cudaFuncGetAttributes(&attrs, kernel) != cudaSuccess) + return false; + if (non_portable + && cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) != cudaSuccess) + return false; + return cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + max_optin - static_cast(attrs.sharedSizeBytes)) + == cudaSuccess; + }; + TLLM_CHECK_WITH_INFO(configure(m3_router_tc_kernel<8>, false), "failed to configure the router kernel"); + if (!configure(m3_router_tc_kernel<16>, true)) + { + cudaGetLastError(); + return 8; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(8 * 16); + cfg.blockDim = dim3(256); + cfg.dynamicSmemBytes + = static_cast(kTcExperts) * (ks + 4) * 4 + static_cast(kTcMaxM) * (ks + 8) * 2; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeClusterDimension; + attr[0].val.clusterDim.x = 16; + attr[0].val.clusterDim.y = 1; + attr[0].val.clusterDim.z = 1; + cfg.attrs = attr; + cfg.numAttrs = 1; + int clusters = 0; + if (cudaOccupancyMaxActiveClusters(&clusters, m3_router_tc_kernel<16>, &cfg) != cudaSuccess) + { + cudaGetLastError(); + return 8; + } + return clusters >= 8 ? 16 : 8; + }(); + 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) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.h b/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.h new file mode 100644 index 000000000000..6f4467883b24 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.h @@ -0,0 +1,46 @@ +/* + * 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 +{ + +//! Low-M fp32 router GEMV: `out[M, E] = float(x[M, K]) @ w[E, K]^T` with exact fp32 FMA, one CTA per expert, +//! 1 <= M <= 16, K % 1024 == 0. With `qOut` the CTAs also emit the NVFP4 activation (packed e2m1 [M, K / 2], linear +//! E4M3 scales [M, K / 16]) for global scale `*sfScale`, bit-identical to `fp4_quantize(x, sfScale, 16, False, +//! False)`. +void invokeMinimaxM3RouterGemv(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t k, float const* w, int32_t e, + float* out, uint8_t* qOut, uint8_t* sfOut, float const* sfScale, bool pdl, cudaStream_t stream); + +//! TF32 tensor-core router for 128 experts and 1 <= M <= 48 (K % 1024 == 0): `out[M, 128]` with the weights rounded +//! to TF32 (round to nearest, ties away) and fp32 accumulation, independent of the cluster size picked. Optional +//! NVFP4 activation as in `invokeMinimaxM3RouterGemv`. +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); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.cu b/cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.cu new file mode 100644 index 000000000000..91b4c38d0fda --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.cu @@ -0,0 +1,525 @@ +/* + * 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/minimaxM3/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). +// +// 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 minimaxM3 +{ +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 + +// 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; +} + +// Every CTA of the grid must be co-resident: a row's CTAs wait for each other's records. +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; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + 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; + + float const tempInv = 1.0f / (a.temperatures[row] + kTempEpsilon); + float const topP = a.top_ps[row]; + unsigned long long const seed = static_cast(a.seeds[row]); + unsigned long long const offset = static_cast(a.offsets[row]); + + // ---- 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; + 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) + { + do + v = ld_sys(xr + (q / kSplits * kMaxSplits + q % kSplits) * kEx); + while ((v & ((1ULL << kTagBits) - 1ULL)) != ktag); + s.rk[q] = v >> kTagBits; + } + else + { + do + v = ld_sys(xr + (q / kSplits * kMaxSplits + q % kSplits) * kEx + f); + 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) + 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_ctas = 0; + if (max_ctas == 0) + { + int dev = 0, sms = 0, per_sm = 0; + cudaError_t err = cudaGetDevice(&dev); + if (err == cudaSuccess) + { + err = cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); + } + if (err == cudaSuccess) + { + err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &per_sm, sampled_accept_kernel, kThreads, sizeof(Shared)); + } + if (err != cudaSuccess) + { + return err; + } + max_ctas = per_sm * sms; + } + TLLM_CHECK_WITH_INFO(a.rows * kSplits <= max_ctas, "m3_sampled_accept: %d CTAs cannot all be co-resident (max %d)", + a.rows * kSplits, max_ctas); + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(a.rows * kSplits); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = sizeof(Shared); + cfg.stream = stream; + cudaLaunchAttribute attr; + attr.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr.val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; + cfg.attrs = &attr; + cfg.numAttrs = 1; + 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 launch<8>(a, usePdl, stream); + } + return launch<4>(a, usePdl, stream); +} + +} // namespace sampledAccept +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.h b/cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.h new file mode 100644 index 000000000000..1dc211a6f8a8 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.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 minimaxM3 +{ +//! 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 +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; +}; + +//! 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 minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/selectLong.cu b/cpp/tensorrt_llm/kernels/minimaxM3/selectLong.cu new file mode 100644 index 000000000000..39351f0ad3ce --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/selectLong.cu @@ -0,0 +1,235 @@ +/* + * 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/minimaxM3/selectLong.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +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) +{ + 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); + 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(); + 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_long_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) +{ + __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); + else if (valid <= 4096) + selectRow<4>(sm, rowBase, blockStride, valid, initBlocks, localStart); + else + selectRow(sm, rowBase, blockStride, valid, initBlocks, localStart); + __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 invokeMinimaxM3SelectLong(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, + cudaStream_t stream) +{ + if (numBlocks > kMinimaxM3SelectLongMaxBlocks) + { + 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_long_kernel, scores, static_cast(headStride), + static_cast(blockStride), static_cast(queryStride), nValidBlocks, out, + static_cast(outQueryStride), static_cast(outHeadStride), numKvHeads, numBlocks, + initBlocks, localBlocks, earlyTrigger); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/selectLong.h b/cpp/tensorrt_llm/kernels/minimaxM3/selectLong.h new file mode 100644 index 000000000000..258954235317 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/selectLong.h @@ -0,0 +1,45 @@ +/* + * 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 +{ + +//! Largest row (blocks per query) the long-row top-16 select covers: 9 keys per thread of one 1024-thread CTA. +constexpr int32_t kMinimaxM3SelectLongMaxBlocks = 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]. With earlyTrigger the dependent +//! grid is released right after this grid's dependency wait instead of at its exit. +cudaError_t invokeMinimaxM3SelectLong(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, + cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/shardAccept.cu b/cpp/tensorrt_llm/kernels/minimaxM3/shardAccept.cu new file mode 100644 index 000000000000..afc70e0e2c35 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/shardAccept.cu @@ -0,0 +1,204 @@ +/* + * 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/minimaxM3/shardAccept.h" + +#include +#include +#include + +// Fused greedy Eagle3 verification for the MiniMax-M3 target head (GB300, TP-sharded vocabulary, <= 48 verify rows). +// +// One kernel replaces the all-gather of the target logits, the argmax and the strict-acceptance ops: +// +// * Each logit of this rank's BF16 shard 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 keys order the BF16 values exactly, so the tokens are the argmax of the gathered logits. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3 +{ +namespace shardAccept +{ +namespace +{ +__device__ __forceinline__ unsigned long long make_key(float v, uint32_t index) +{ + uint32_t u = __float_as_uint(v); + // -0.0 and +0.0 compare equal (torch.argmax takes the lower index among them): one key for both. + u = (u == 0x80000000u) ? 0u : u; + 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; +} + +// Greedy verification from this rank's logits shard (verify-row count <= kMaxRows): per-row argmax keys over the +// shard, the last CTA exchanges them across ranks (multicast slots [2][kMaxRows][world], Lamport 0 = empty) and runs +// the acceptance epilogue. +__global__ void __launch_bounds__(kThreads, 1) shard_accept_kernel(__nv_bfloat16 const* __restrict__ logits, int M, + int Vs, long long ld, int vocab_offset, unsigned long long* __restrict__ keys, unsigned int* __restrict__ counters, + unsigned long long* __restrict__ ws_uc, unsigned long long* __restrict__ ws_mc, int rank, int world, AcceptArgs acc) +{ + __shared__ bool s_last; + __shared__ int s_tok[kMaxRows]; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + int const nwarps = blockDim.x >> 5; + int const groups = (Vs + 7) / 8; + int const g0 = static_cast((static_cast(blockIdx.x) * groups) / gridDim.x); + int const g1 = static_cast((static_cast(blockIdx.x + 1) * groups) / gridDim.x); + bool const vec = (ld % 8) == 0; + for (int r = warp; r < M; r += nwarps) + { + unsigned long long best = 0ull; + __nv_bfloat16 const* row = logits + r * ld; + for (int gi = g0 + lane; gi < g1; gi += 32) + { + int const c = gi * 8; + if (vec && c + 8 <= Vs) + { + uint4 const raw = *reinterpret_cast(row + c); + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); +#pragma unroll + for (int q = 0; q < 8; ++q) + { + unsigned long long const k + = make_key(__bfloat162float(e[q]), static_cast(vocab_offset + c + q)); + best = k > best ? k : best; + } + } + else + { + for (int q = 0; q < 8 && c + q < Vs; ++q) + { + unsigned long long const k + = make_key(__bfloat162float(row[c + q]), static_cast(vocab_offset + c + q)); + best = k > best ? k : best; + } + } + } +#pragma unroll + for (int o = 1; o < 32; o <<= 1) + best = shfl_max(best, o); + if (lane == 0 && best != 0ull) + atomicMax(keys + r, best); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) + { + unsigned int const prev = atomicAdd(counters, 1u); + s_last = prev == gridDim.x - 1; + } + __syncthreads(); + if (!s_last) + return; + __threadfence(); + unsigned int const call = *reinterpret_cast(counters + 1); + int const half = static_cast(call & 1u); + for (int t = threadIdx.x; t < M; t += blockDim.x) + { + unsigned long long const mine = atomicExch(keys + t, 0ull); + unsigned long long* slot_mc = ws_mc + (half * kMaxRows + t) * world; + unsigned long long* slot_uc = ws_uc + (half * kMaxRows + t) * world; + slot_mc[rank] = mine; + unsigned long long best = 0ull; + for (int r = 0; r < 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)); + } + __syncthreads(); + if (threadIdx.x == 0) + { + counters[0] = 0u; + counters[1] = call + 1u; + } + if (threadIdx.x < acc.nreq) + { + int const b = threadIdx.x; + int const k1 = acc.draft_len + 1; + int n = 1; + bool run = true; + for (int k = 0; k < acc.draft_len; ++k) + { + int const tt = s_tok[b * k1 + k]; + acc.acc_tokens[b * k1 + k] = tt; + run = run && acc.draft[b * acc.draft_len + k] == tt; + n += run ? 1 : 0; + } + acc.acc_tokens[b * k1 + acc.draft_len] = s_tok[b * k1 + acc.draft_len]; + if (acc.force_vec != nullptr) + n = acc.force_vec[b]; + else if (acc.force_total > 0) + n = acc.force_total; + acc.num_acc[b] = n; + acc.gather_ids[b] = static_cast(acc.batch_idx[b]) * k1 + n - 1; + if (acc.next_new != nullptr) + acc.next_new[b * acc.next_stride] = s_tok[acc.batch_idx[b] * k1 + n - 1]; + } +} + +} // namespace + +cudaError_t invokeShardAccept(__nv_bfloat16 const* logits, int M, int Vs, long long ld, int vocabOffset, + unsigned long long* keys, unsigned int* counters, unsigned long long* wsUc, unsigned long long* wsMc, int rank, + int world, AcceptArgs const& acc, int numCtas, cudaStream_t stream) +{ + cudaLaunchAttribute attr[1]; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(static_cast(numCtas)); + cfg.blockDim = dim3(kThreads); + cfg.stream = stream; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = 1; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx( + &cfg, shard_accept_kernel, logits, M, Vs, ld, vocabOffset, keys, counters, wsUc, wsMc, rank, world, acc); +} + +} // namespace shardAccept +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/shardAccept.h b/cpp/tensorrt_llm/kernels/minimaxM3/shardAccept.h new file mode 100644 index 000000000000..e41982b9d357 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/shardAccept.h @@ -0,0 +1,67 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3 +{ +//! Fused greedy Eagle3 verification of the MiniMax-M3 NVFP4 TP4 decode path (sm_103a): the cross-rank argmax of the +//! target logits shard and the strict acceptance in one kernel. See shardAccept.cu. +namespace shardAccept +{ + +constexpr int kMaxRows = 48; // verify rows of one launch +constexpr int kThreads = 512; // threads per CTA + +// Greedy-acceptance epilogue run by the last CTA once the global per-row argmax is known (rows = nreq requests x +// (draft_len + 1) verify positions, request-major). Mirrors the strict path of SpecWorkerBase: accepted = the target +// tokens, count = 1 + length of the matching draft prefix, optionally forced; plus the first draft pass's row index. +struct AcceptArgs +{ + int const* draft; // [nreq, draft_len] previous draft tokens + int const* batch_idx; // [nreq] request row in the verify layout + int* acc_tokens; // [nreq, draft_len + 1] + int* num_acc; // [nreq] + long long* gather_ids; // [nreq] last accepted row + int draft_len; + int nreq; + int force_total; // > 0 overrides the accepted count + int* next_new; // optional [nreq, next_stride]: column 0 = last accepted token (next step's input) + int next_stride; + int const* force_vec; // optional [nreq] per-request accepted count; overrides force_total +}; + +//! Greedy verification from this rank's BF16 logits shard [M, Vs] (row stride ld), M <= kMaxRows, over numCtas +//! kThreads-thread CTAs, followed by the strict acceptance epilogue. +cudaError_t invokeShardAccept(__nv_bfloat16 const* logits, int M, int Vs, long long ld, int vocabOffset, + unsigned long long* keys, unsigned int* counters, unsigned long long* wsUc, unsigned long long* wsMc, int rank, + int world, AcceptArgs const& acc, int numCtas, cudaStream_t stream); + +} // namespace shardAccept +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/sparseDecodeLc.cu b/cpp/tensorrt_llm/kernels/minimaxM3/sparseDecodeLc.cu new file mode 100644 index 000000000000..e01fa78927a5 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/sparseDecodeLc.cu @@ -0,0 +1,561 @@ +/* + * 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. + +#include "tensorrt_llm/kernels/minimaxM3/sparseDecodeLc.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(); + +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]; + 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 +__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 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"); + + // 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 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). + 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 + 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)); + } + } + + asm volatile("griddepcontrol.wait;" ::: "memory"); + if constexpr (EARLY) + { + asm volatile("griddepcontrol.launch_dependents;"); + } + int blk = static_cast(rank) < real_topk ? p.topk[tok * p.t_tok + rank * p.t_k] : -1; + bool const active = blk >= 0; + int64_t const page_pw = !active ? 0 : 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(); + 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. + int const h = static_cast(rank); + float lse_c[kTopk]; + float lse_max = kNegInf; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + { + lse_c[c] = s.recv_lse[c]; + lse_max = fmaxf(lse_max, lse_c[c]); + } + lse_max = lse_max == kNegInf ? 0.F : lse_max; + float denom = 0.F; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + { + lse_c[c] = ex2(lse_c[c] - lse_max); + denom += lse_c[c]; + } + denom = denom > 0.F ? denom : 1.F; + float o = 0.F; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + o += s.recv[c][tid] * __fdiv_rn(lse_c[c], denom); + __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, sizeof(Smem)); + 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 = sizeof(Smem); + 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 invokeMinimaxM3SparseDecodeLc(MinimaxM3SparseDecodeLcParams const& params, int32_t numTokens, + bool linearSoftmax, bool mxfp8, bool earlyTrigger, bool usePdl, cudaStream_t stream) +{ + if (!earlyTrigger) + { + return cudaErrorNotSupported; + } + 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/minimaxM3/sparseDecodeLc.h b/cpp/tensorrt_llm/kernels/minimaxM3/sparseDecodeLc.h new file mode 100644 index 000000000000..659531f725bf --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/sparseDecodeLc.h @@ -0,0 +1,70 @@ +/* + * 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. 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; +}; + +//! 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 invokeMinimaxM3SparseDecodeLc(MinimaxM3SparseDecodeLcParams const& params, int32_t numTokens, + bool linearSoftmax, bool mxfp8, bool earlyTrigger, bool usePdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu b/cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu new file mode 100644 index 000000000000..4069da14faf0 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu @@ -0,0 +1,425 @@ +/* + * 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/minimaxM3/specTail.h" + +#include +#include +#include + +// 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 with +// the in-tree RMSNorm arithmetic (fp32 statistic, (x * rsqrt) * w, one bf16 rounding); step_gather / slot_scatter move +// the step's integer outputs between the spec sampler's store and the engine's input buffers. +// +// draft_input replaces a fused draft pass's embedding lookup, its two RMSNorms (input_layernorm on the embedding, +// hidden_norm on the carried hidden state, the same arithmetic) and their concatenation with one launch. It optionally +// gathers the carried hidden rows (the second pass reads the last accepted row of each request) and writes them out +// as the residual of the decoder layer. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3 +{ +namespace specTail +{ +namespace +{ +constexpr int kVec = 8; // bf16 per 16-byte access + +__device__ __forceinline__ float block_sum(float v, float* red) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + int const warp = threadIdx.x >> 5; + if ((threadIdx.x & 31) == 0) + red[warp] = v; + __syncthreads(); + float s = 0.f; +#pragma unroll + for (int w = 0; w < kThreads / 32; ++w) + s += red[w]; + return s; +} + +// blockIdx.x = row, blockIdx.y = 0 (embedding half) / 1 (hidden half). H % (kThreads * kVec) == 0 is not required; +// H % kVec == 0 and H <= kMaxPerThread * kThreads * kVec are. +template +__global__ void __launch_bounds__(kThreads) + draft_input_kernel(int const* __restrict__ ids, int vocab, __nv_bfloat16 const* __restrict__ table, + __nv_bfloat16 const* __restrict__ hidden, long long const* __restrict__ hidden_rows, + __nv_bfloat16 const* __restrict__ w_embed, __nv_bfloat16 const* __restrict__ w_hidden, float eps_embed, + float eps_hidden, int H, __nv_bfloat16* __restrict__ out, __nv_bfloat16* __restrict__ residual) +{ + __shared__ float red[kThreads / 32]; + int const row = blockIdx.x; + bool const is_hidden = blockIdx.y == 1; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + __nv_bfloat16 const* src; + bool valid = true; + if (is_hidden) + { + long long const r = hidden_rows ? hidden_rows[row] : row; + src = hidden + r * H; + } + else + { + int const id = ids[row]; + valid = id >= 0 && id < vocab; + src = table + static_cast(valid ? id : 0) * H; + } + float v[kIters][kVec]; + float ss = 0.f; +#pragma unroll + for (int it = 0; it < kIters; ++it) + { + int const c = (it * kThreads + threadIdx.x) * kVec; + if (c < H) + { + uint4 const raw = *reinterpret_cast(src + c); + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); +#pragma unroll + for (int j = 0; j < kVec; ++j) + { + v[it][j] = valid ? __bfloat162float(e[j]) : 0.f; + ss += v[it][j] * v[it][j]; + } + if (is_hidden && residual) + *reinterpret_cast(residual + static_cast(row) * H + c) = raw; + } + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + 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) + { + 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); + uint4 o; + __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); +#pragma unroll + for (int j = 0; j < kVec; ++j) + oe[j] = __float2bfloat16(v[it][j] * rs * __bfloat162float(we[j])); + *reinterpret_cast(dst + c) = o; + } + } +} + +// blockIdx.x = row, blockIdx.y = chunk c: dst[row, c*H:(c+1)*H] = rmsnorm(src[row, c*H:(c+1)*H]) * w[c]. +template +__global__ void __launch_bounds__(kThreads) chunk_norm_kernel(__nv_bfloat16 const* __restrict__ src, + long long src_stride, __nv_bfloat16 const* __restrict__ w, float eps, int H, int C, __nv_bfloat16* __restrict__ dst) +{ + __shared__ float red[kThreads / 32]; + int const row = blockIdx.x; + int const c = blockIdx.y; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + __nv_bfloat16 const* in = src + row * src_stride + static_cast(c) * H; + float v[kIters][kVec]; + float ss = 0.f; +#pragma unroll + for (int it = 0; it < kIters; ++it) + { + int const col = (it * kThreads + threadIdx.x) * kVec; + if (col < H) + { + uint4 const raw = *reinterpret_cast(in + col); + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); +#pragma unroll + for (int j = 0; j < kVec; ++j) + { + v[it][j] = __bfloat162float(e[j]); + ss += v[it][j] * v[it][j]; + } + } + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float const rs = rsqrtf(block_sum(ss, red) / H + eps); + __nv_bfloat16 const* wc = w + static_cast(c) * H; + __nv_bfloat16* out = dst + (static_cast(row) * C + c) * H; +#pragma unroll + for (int it = 0; it < kIters; ++it) + { + int const col = (it * kThreads + threadIdx.x) * kVec; + if (col < H) + { + uint4 const wr = *reinterpret_cast(wc + col); + __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr); + uint4 o; + __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); +#pragma unroll + for (int j = 0; j < kVec; ++j) + oe[j] = __float2bfloat16(v[it][j] * rs * __bfloat162float(we[j])); + *reinterpret_cast(out + col) = o; + } + } +} + +// Eagle3 hidden-state capture with the drafter's fc_norm folded in (one row per CTA): +// dst[row, :H] = rmsnorm(bf16(h[row] + r[row])) * w (r optional) +// with chunk_norm's arithmetic, so fc reads the same values the capture + chunk_norm pair produced. +template +__global__ void __launch_bounds__(kThreads) capture_norm_kernel(__nv_bfloat16 const* __restrict__ h, long long h_stride, + __nv_bfloat16 const* __restrict__ r, long long r_stride, __nv_bfloat16 const* __restrict__ w, float eps, int H, + __nv_bfloat16* __restrict__ dst, long long dst_stride) +{ + __shared__ float red[kThreads / 32]; + int const row = blockIdx.x; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + float v[kIters][kVec]; + float ss = 0.f; +#pragma unroll + for (int it = 0; it < kIters; ++it) + { + int const col = (it * kThreads + threadIdx.x) * kVec; + if (col < H) + { + uint4 const raw = *reinterpret_cast(h + row * h_stride + col); + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); + uint4 rraw = make_uint4(0u, 0u, 0u, 0u); + if (r != nullptr) + rraw = *reinterpret_cast(r + row * r_stride + col); + __nv_bfloat16 const* re = reinterpret_cast<__nv_bfloat16 const*>(&rraw); +#pragma unroll + for (int j = 0; j < kVec; ++j) + { + float const x = __bfloat162float(e[j]); + v[it][j] = r != nullptr ? __bfloat162float(__float2bfloat16(x + __bfloat162float(re[j]))) : x; + ss += v[it][j] * v[it][j]; + } + } + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float const rs = rsqrtf(block_sum(ss, red) / H + eps); + __nv_bfloat16* out = dst + row * dst_stride; +#pragma unroll + for (int it = 0; it < kIters; ++it) + { + int const col = (it * kThreads + threadIdx.x) * kVec; + if (col < H) + { + uint4 const wr = *reinterpret_cast(w + col); + __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr); + uint4 o; + __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); +#pragma unroll + for (int j = 0; j < kVec; ++j) + oe[j] = __float2bfloat16(v[it][j] * rs * __bfloat162float(we[j])); + *reinterpret_cast(out + col) = o; + } + } +} + +// Overlap-scheduler step inputs from the previous step's device outputs (one launch for the engine's gathers): +// ids[ids0 + p*T + t] = new_tok[t * nt_s0 + slot_p * nt_s1] +// draft[d0 + p*D + d] = next_draft[slot_p * nd_s0 + d * nd_s1] +// pos_off[i] = i in [b0*T, (b0+P)*T) ? lens[ppi[i - b0*T]] : 0 (i < n_pos) +// kv_off[i] = i in [b0, b0+P) ? lens[slot_{i-b0}] - T : 0 (i < n_kv) +__global__ void step_gather_kernel(int const* __restrict__ new_tok, long long nt_s0, long long nt_s1, + int const* __restrict__ next_draft, long long nd_s0, long long nd_s1, int const* __restrict__ lens, + int const* __restrict__ slots, int const* __restrict__ ppi, int P, int T, int D, int* __restrict__ ids, int ids0, + int* __restrict__ draft, int d0, int* __restrict__ pos_off, int n_pos, int* __restrict__ kv_off, int n_kv, int b0) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + int const tid = blockIdx.x * blockDim.x + threadIdx.x; + int const stride = gridDim.x * blockDim.x; + for (int i = tid; i < P * T; i += stride) + { + int const p = i / T, t = i - p * T; + ids[ids0 + i] = new_tok[t * nt_s0 + static_cast(slots[p]) * nt_s1]; + } + for (int i = tid; i < P * D; i += stride) + { + int const p = i / D, d = i - p * D; + draft[d0 + i] = next_draft[static_cast(slots[p]) * nd_s0 + d * nd_s1]; + } + for (int i = tid; i < n_pos; i += stride) + { + int const k = i - b0 * T; + pos_off[i] = (k >= 0 && k < P * T) ? lens[ppi[k]] : 0; + } + for (int i = tid; i < n_kv; i += stride) + { + int const k = i - b0; + kv_off[i] = (k >= 0 && k < P) ? lens[slots[k]] - T : 0; + } +} + +// Sampler stores from one step's outputs (one launch for its index_copy_s), zero-padding to the store widths: +// a[slot_i][c] for c < wa (from src_a[i][c] if c < sa_w), b likewise, lens[slot_i] = src_lens[i], d likewise. +__global__ void slot_scatter_kernel(long long const* __restrict__ slots, int n, int const* __restrict__ sa, + long long sa_s0, long long sa_s1, int sa_w, int* __restrict__ da, long long da_s0, long long da_s1, int wa, + int const* __restrict__ sb, long long sb_s0, long long sb_s1, int sb_w, int* __restrict__ db, long long db_s0, + long long db_s1, int wb, int const* __restrict__ sl, int* __restrict__ dl, int const* __restrict__ sd, + long long sd_s0, long long sd_s1, int sd_w, int* __restrict__ dd, long long dd_s0, long long dd_s1, int wd) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + int const i = blockIdx.x; + if (i >= n) + return; + long long const s = slots[i]; + for (int c = threadIdx.x; c < wa; c += blockDim.x) + da[s * da_s0 + c * da_s1] = c < sa_w ? sa[i * sa_s0 + c * sa_s1] : 0; + for (int c = threadIdx.x; c < wb; c += blockDim.x) + db[s * db_s0 + c * db_s1] = c < sb_w ? sb[i * sb_s0 + c * sb_s1] : 0; + for (int c = threadIdx.x; c < wd; c += blockDim.x) + dd[s * dd_s0 + c * dd_s1] = c < sd_w ? sd[i * sd_s0 + c * sd_s1] : 0; + if (threadIdx.x == 0) + dl[s] = sl[i]; +} + +cudaLaunchConfig_t small_cfg(int grid, int block, bool pdl, cudaLaunchAttribute* attr, cudaStream_t stream) +{ + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(grid); + cfg.blockDim = dim3(block); + cfg.stream = stream; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cfg; +} + +int norm_iters(int H) +{ + return (H / kVec + kThreads - 1) / kThreads; +} + +} // namespace + +cudaError_t invokeDraftInput(int const* ids, int vocab, __nv_bfloat16 const* table, __nv_bfloat16 const* hidden, + long long const* rows, __nv_bfloat16 const* wEmbed, __nv_bfloat16 const* wHidden, float epsEmbed, float epsHidden, + int H, int M, __nv_bfloat16* out, __nv_bfloat16* residual, bool pdl, cudaStream_t stream) +{ + int const iters = norm_iters(H); + cudaLaunchAttribute attr[1]; + cudaLaunchConfig_t cfg = small_cfg(1, kThreads, pdl, attr, stream); + cfg.gridDim = dim3(M, 2); +#define LAUNCH(I) \ + cudaLaunchKernelEx(&cfg, draft_input_kernel, ids, vocab, table, hidden, rows, wEmbed, wHidden, epsEmbed, \ + epsHidden, H, out, residual) + switch (iters) + { + case 0: + case 1: return LAUNCH(1); + case 2: return LAUNCH(2); + case 3: return LAUNCH(3); + case 4: return LAUNCH(4); + default: return cudaErrorInvalidValue; + } +#undef LAUNCH +} + +cudaError_t invokeChunkNorm(__nv_bfloat16 const* src, long long srcStride, __nv_bfloat16 const* w, float eps, int H, + int C, int M, __nv_bfloat16* dst, bool pdl, cudaStream_t stream) +{ + int const iters = norm_iters(H); + cudaLaunchAttribute attr[1]; + cudaLaunchConfig_t cfg = small_cfg(1, kThreads, pdl, attr, stream); + cfg.gridDim = dim3(M, C); +#define LAUNCH(I) cudaLaunchKernelEx(&cfg, chunk_norm_kernel, src, srcStride, w, eps, H, C, dst) + switch (iters) + { + case 0: + case 1: return LAUNCH(1); + case 2: return LAUNCH(2); + case 3: return LAUNCH(3); + case 4: return LAUNCH(4); + default: return cudaErrorInvalidValue; + } +#undef LAUNCH +} + +cudaError_t invokeCaptureNorm(__nv_bfloat16 const* h, long long hStride, __nv_bfloat16 const* res, long long resStride, + __nv_bfloat16 const* w, float eps, int H, int M, __nv_bfloat16* dst, long long dstStride, cudaStream_t stream) +{ + if (M == 0) + { + return cudaSuccess; + } + int const iters = norm_iters(H); + cudaLaunchAttribute attr[1]; + cudaLaunchConfig_t cfg = small_cfg(M, kThreads, true, attr, stream); +#define LAUNCH(I) \ + cudaLaunchKernelEx(&cfg, capture_norm_kernel, h, hStride, res, resStride, w, eps, H, dst, dstStride) + switch (iters) + { + case 0: + case 1: return LAUNCH(1); + case 2: return LAUNCH(2); + case 3: return LAUNCH(3); + case 4: return LAUNCH(4); + default: return cudaErrorInvalidValue; + } +#undef LAUNCH +} + +cudaError_t invokeStepGather(int const* newTok, long long ntS0, long long ntS1, int const* nextDraft, long long ndS0, + long long ndS1, int const* lens, int const* slots, int const* ppi, int P, int T, int D, int* ids, int ids0, + int* draft, int d0, int* posOff, int nPos, int* kvOff, int nKv, int b0, int work, cudaStream_t stream) +{ + cudaLaunchAttribute attr[1]; + cudaLaunchConfig_t cfg = small_cfg(std::max(1, std::min(64, (work + 255) / 256)), 256, true, attr, stream); + return cudaLaunchKernelEx(&cfg, step_gather_kernel, newTok, ntS0, ntS1, nextDraft, ndS0, ndS1, lens, slots, ppi, P, + T, D, ids, ids0, draft, d0, posOff, nPos, kvOff, nKv, b0); +} + +cudaError_t invokeSlotScatter(long long const* slots, int n, IntView sa, IntView da, IntView sb, IntView db, + int const* sl, int* dl, IntView sd, IntView dd, cudaStream_t stream) +{ + if (n == 0) + { + return cudaSuccess; + } + cudaLaunchAttribute attr[1]; + cudaLaunchConfig_t cfg = small_cfg(n, 32, true, attr, stream); + return cudaLaunchKernelEx(&cfg, slot_scatter_kernel, slots, n, sa.data, sa.s0, sa.s1, sa.width, da.data, da.s0, + da.s1, da.width, sb.data, sb.s0, sb.s1, sb.width, db.data, db.s0, db.s1, db.width, sl, dl, sd.data, sd.s0, + sd.s1, sd.width, dd.data, dd.s0, dd.s1, dd.width); +} + +} // namespace specTail +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/specTail.h b/cpp/tensorrt_llm/kernels/minimaxM3/specTail.h new file mode 100644 index 000000000000..c5bdf75ff482 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/specTail.h @@ -0,0 +1,80 @@ +/* + * 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 minimaxM3 +{ +//! Small fused kernels of the MiniMax-M3 Eagle3 step: the drafter's fc_norm RMSNorms over the captured hidden states +//! and the step's input / output bookkeeping (GB300, decode). All launches use programmatic dependent launch where +//! noted. +namespace specTail +{ + +//! Threads per CTA of the norm kernels, and the hidden-size bound of their register tiles: +//! H % 8 == 0 and H <= kMaxIters * kThreads * 8. +constexpr int kThreads = 256; +constexpr int kMaxIters = 4; + +//! out[r] = [rmsnorm(table[ids[r]]) * wEmbed, rmsnorm(hidden[rows ? rows[r] : r]) * wHidden] ([M, 2H] BF16; ids +//! outside [0, vocab) embed to zero); residual[r] = hidden[rows[r]] when residual is set. +cudaError_t invokeDraftInput(int const* ids, int vocab, __nv_bfloat16 const* table, __nv_bfloat16 const* hidden, + long long const* rows, __nv_bfloat16 const* wEmbed, __nv_bfloat16 const* wHidden, float epsEmbed, float epsHidden, + int H, int M, __nv_bfloat16* out, __nv_bfloat16* residual, bool pdl, cudaStream_t stream); + +//! dst[r, c*H:(c+1)*H] = rmsnorm(src[r, c*H:(c+1)*H]) * w[c] for c < C (src row stride srcStride), dst [M, C*H]. +cudaError_t invokeChunkNorm(__nv_bfloat16 const* src, long long srcStride, __nv_bfloat16 const* w, float eps, int H, + int C, int M, __nv_bfloat16* dst, bool pdl, cudaStream_t stream); + +//! dst[r, :H] = rmsnorm(bf16(h[r] + res[r])) * w (res optional), with chunk-norm rounding. +cudaError_t invokeCaptureNorm(__nv_bfloat16 const* h, long long hStride, __nv_bfloat16 const* res, long long resStride, + __nv_bfloat16 const* w, float eps, int H, int M, __nv_bfloat16* dst, long long dstStride, cudaStream_t stream); + +//! The overlap scheduler's per-step input gathers (see step_gather_kernel); work = the largest of its four loops. +cudaError_t invokeStepGather(int const* newTok, long long ntS0, long long ntS1, int const* nextDraft, long long ndS0, + long long ndS1, int const* lens, int const* slots, int const* ppi, int P, int T, int D, int* ids, int ids0, + int* draft, int d0, int* posOff, int nPos, int* kvOff, int nKv, int b0, int work, cudaStream_t stream); + +//! A strided int32 2-D view: element (i, c) at data[i * s0 + c * s1], `width` columns. +struct IntView +{ + int* data; + long long s0; + long long s1; + int width; +}; + +//! The spec sampler's slot stores (see slot_scatter_kernel): for i < n, da / db / dd rows slots[i] get sa / sb / sd +//! row i, zero-padded to the destination widths, and dl[slots[i]] = sl[i]. +cudaError_t invokeSlotScatter(long long const* slots, int n, IntView sa, IntView da, IntView sb, IntView db, + int const* sl, int* dl, IntView sd, IntView dd, cudaStream_t stream); + +} // namespace specTail +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/targetHead.cu b/cpp/tensorrt_llm/kernels/minimaxM3/targetHead.cu new file mode 100644 index 000000000000..847c59779a6e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/targetHead.cu @@ -0,0 +1,372 @@ +/* + * 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/minimaxM3/targetHead.h" + +#include +#include +#include + +// The MiniMax-M3 target lm_head shard GEMV of a verify step of at most 8 rows (GB300, TP4), at stock precision: BF16 +// weight and activation, mma.m16n8k16 BF16 with FP32 accumulation over the full K, one rounding to BF16 per logit +// (the cuBLAS head differs only in the FP32 summation order). The full logit row is produced. +// +// * Lossless 12-bit weight 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; the decoded BF16 bits are the checkpoint's. Code 15 marks an exception: it decodes to +0 (it +// adds nothing to the MMA sums), and the segment epilogue adds the exact product x * w from a per-tile list (a +// product of two BF16 values is exact in FP32), so every product enters the sum exactly as in the BF16 head. +// * Layout: the weight is permuted into mma A-fragment order (per 16-row tile and k16 step, lane (g, j) holds +// {row g k 2j..2j+1, row g+8 k 2j.., row g k 2j+8.., row g+8 k 2j+8..}); per unit of 4 k16 steps a lane reads +// 48 bytes: the sign/mantissa bytes of steps 0-1 and of steps 2-3, then 4 code words (even values' nibbles in the +// low half, odd values' in the high half, so the codes of value pair p are (word >> 4p) & 0x000f000f). +// * A persistent grid (one 16-warp CTA per SM) splits the (tile, unit) space into equal contiguous ranges, one per +// warp, each streamed through a 3-unit cp.async ring in shared memory. The first units are issued before the PDL +// grid-dependency wait (the weight does not depend on the previous kernel); the activation is staged in shared +// memory in B-fragment order after it. +// * A tile split between warps writes per-warp FP32 partials; the last finisher (per-tile counter) adds them in warp +// order, so the result is deterministic run to run. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3 +{ +namespace targetHead +{ +namespace +{ + +constexpr int kUnit = 4; // k16 steps per coded unit + +struct Args +{ + Params p; + int tiles; + int ks; // K / 16 + long long units; // tiles * ks / kUnit +}; + +__device__ __forceinline__ void cp_async16(void* smem, void const* gmem) +{ + uint32_t const sa = static_cast(__cvta_generic_to_shared(smem)); + asm volatile("cp.async.cg.shared.global.L2::256B [%0], [%1], 16;\n" ::"r"(sa), "l"(gmem) : "memory"); +} + +__device__ __forceinline__ void cp_async_commit() +{ + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} + +template +__device__ __forceinline__ void cp_async_wait() +{ + asm volatile("cp.async.wait_group %0;\n" ::"n"(N) : "memory"); +} + +__device__ __forceinline__ void mma_bf16(float (&d)[4], uint4 const& a, uint2 const& b) +{ + 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"(a.x), "r"(a.y), "r"(a.z), "r"(a.w), "r"(b.x), "r"(b.y)); +} + +// Accumulator layout of one tile: c[0..1] = (row g, tokens 2j, 2j + 1), c[2..3] = (row g + 8, ...). +__device__ __forceinline__ void store_tile(Args const& a, int tile, int lane, float const (&c)[4]) +{ + int const g = lane >> 2, j = lane & 3; +#pragma unroll + for (int e = 0; e < 4; ++e) + { + int const n = 2 * j + (e & 1); + int const row = tile * 16 + g + (e >> 1) * 8; + if (n < a.p.M && row < a.p.N) + a.p.out[n * a.p.out_stride + row] = __float2bfloat16_rn(c[e]); + } +} + +// Flush the warp's sums of `tile` over k16 steps [k0, k1). +__device__ __forceinline__ void flush(Args const& a, int warp, int tile, int k0, int k1, int lane, float const (&c)[4]) +{ + if (k0 == 0 && k1 == a.ks) + { + store_tile(a, tile, lane, c); + return; + } + // Warps owning parts of this tile: the owner of its first unit up to the owner of its last one, where + // owner(u) = the largest w with w * units / W <= u. The first owner stores into its slot 1 (the tile is its last + // segment), every later owner into slot 0 (its first segment), so a warp never reuses a slot. + long long const W = static_cast(gridDim.x) * (blockDim.x >> 5); + long long const first = static_cast(tile) * (a.ks / kUnit); + auto owner = [&](long long u) + { + long long w = (u * W) / a.units; + while (w + 1 < W && ((w + 1) * a.units) / W <= u) + ++w; + while (w > 0 && (w * a.units) / W > u) + --w; + return static_cast(w); + }; + int const wf = owner(first); + int const wl = owner(first + a.ks / kUnit - 1); + constexpr int kVals = 4 * 32; + float* mine = a.p.part + (static_cast(warp) * 2 + (warp == wf ? 1 : 0)) * kVals; +#pragma unroll + for (int e = 0; e < 4; ++e) + mine[e * 32 + lane] = c[e]; + __threadfence(); + __syncwarp(); + int done = 0; + if (lane == 0) + done = atomicAdd(&a.p.counter[tile], 1) == wl - wf; + done = __shfl_sync(0xffffffffu, done, 0); + if (!done) + return; + __threadfence(); + float s[4] = {0.f, 0.f, 0.f, 0.f}; + for (int w = wf; w <= wl; ++w) + { + float const* q = a.p.part + (static_cast(w) * 2 + (w == wf ? 1 : 0)) * kVals; +#pragma unroll + for (int e = 0; e < 4; ++e) + s[e] += __ldcg(q + e * 32 + lane); + } + store_tile(a, tile, lane, s); + if (lane == 0) + a.p.counter[tile] = 0; +} + +// 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 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); +} + +// WARPS warps per CTA, each with a cp.async ring of S units (1536 bytes each). The range is walked tile segment by +// tile segment, so the inner loop has no boundary checks; k16 steps alternate between two accumulators (added at the +// segment end). x is staged as [ks / 4][32][4] B fragments, two 16-byte shared loads per unit. +template +__global__ void __launch_bounds__(WARPS * 32, 1) coded_kernel(Args const a) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ uint4 smem4[]; + int const lane = threadIdx.x & 31; + int const wid = threadIdx.x >> 5; + int const warp = blockIdx.x * WARPS + wid; + long long const W = static_cast(gridDim.x) * WARPS; + long long const u0 = warp * a.units / W; + long long const u1 = (warp + 1) * a.units / W; + int const n_units = static_cast(u1 - u0); + int const upt = a.ks / kUnit; // units per tile + uint4 const* wp = a.p.w + u0 * 96 + lane; + uint4* const xs4 = smem4; // [upt][32][2] uint4 + uint4* const ring = smem4 + upt * 64 + wid * (S * 96) + lane; + +#pragma unroll + for (int u = 0; u < S; ++u) + { + if (u < n_units) +#pragma unroll + for (int l = 0; l < 3; ++l) + cp_async16(ring + (u * 3 + l) * 32, wp + (u * 3 + l) * 32); + cp_async_commit(); + } + asm volatile("griddepcontrol.wait;" ::: "memory"); + { + // x[n][8q..8q+7] -> word h = q & 1 of entries (s = q / 2, lane = n * 4 + j), entry index (s / 4, lane, s % 4). + uint32_t* xs32 = reinterpret_cast(xs4); + int const qn = a.ks * 2; + int const items = kMaxRows * qn; + constexpr int kSU = 8; + for (int i0 = threadIdx.x; i0 < items; i0 += kSU * blockDim.x) + { + uint4 v[kSU]; +#pragma unroll + for (int u = 0; u < kSU; ++u) + { + int const i = i0 + u * blockDim.x; + int const n = i / qn, q = i - n * qn; + v[u] = (i < items && n < a.p.M) ? *reinterpret_cast(a.p.x + n * a.p.x_stride + q * 8) + : make_uint4(0u, 0u, 0u, 0u); + } +#pragma unroll + for (int u = 0; u < kSU; ++u) + { + int const i = i0 + u * blockDim.x; + if (i >= items) + continue; + int const n = i / qn, q = i - n * qn; + int const sstep = q >> 1; + int const base = (((sstep >> 2) * 32 + n * 4) * 4 + (sstep & 3)) * 2 + (q & 1); + xs32[base + 0] = v[u].x; + xs32[base + 8] = v[u].y; + xs32[base + 16] = v[u].z; + xs32[base + 24] = v[u].w; + } + } + } + __syncthreads(); + asm volatile("griddepcontrol.launch_dependents;"); + + uint4 const* xl = xs4 + lane * 2; + int tile = static_cast(u0 / upt); + int ku = static_cast(u0 % upt); // unit within the tile + int done = 0, slot = 0; + while (done < n_units) + { + int const seg = min(n_units - done, upt - ku); + float ca[4] = {0.f, 0.f, 0.f, 0.f}; + float cb[4] = {0.f, 0.f, 0.f, 0.f}; + int const e0 = __ldg(a.p.exc_off + tile), e1 = __ldg(a.p.exc_off + tile + 1); + for (int i = 0; i < seg; ++i) + { + cp_async_wait(); + uint4 const sm01 = ring[slot * 96]; + uint4 const sm23 = ring[slot * 96 + 32]; + uint4 const cw = ring[slot * 96 + 64]; + uint4 const xa = xl[(ku + i) * 64]; + uint4 const xb = xl[(ku + i) * 64 + 1]; + 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}; + uint2 const xb2[4] + = {make_uint2(xa.x, xa.y), make_uint2(xa.z, xa.w), make_uint2(xb.x, xb.y), make_uint2(xb.z, xb.w)}; +#pragma unroll + for (int ci = 0; ci < 4; ++ci) + { + uint4 wv; + wv.x = decode_pair(smw[2 * ci], 0, cws[ci], 0, a.p.exp_base); + wv.y = decode_pair(smw[2 * ci], 2, cws[ci], 1, a.p.exp_base); + wv.z = decode_pair(smw[2 * ci + 1], 0, cws[ci], 2, a.p.exp_base); + wv.w = decode_pair(smw[2 * ci + 1], 2, cws[ci], 3, a.p.exp_base); + if (ci & 1) + mma_bf16(cb, wv, xb2[ci]); + else + mma_bf16(ca, wv, xb2[ci]); + } + // Refill the slot only after its values were consumed (no ordering between ld.shared and cp.async). + int const nxt = done + i + S; + if (nxt < n_units) +#pragma unroll + for (int l = 0; l < 3; ++l) + cp_async16(ring + (slot * 3 + l) * 32, wp + (nxt * 3 + l) * 32); + cp_async_commit(); + slot = slot + 1 == S ? 0 : slot + 1; + } +#pragma unroll + for (int e = 0; e < 4; ++e) + ca[e] += cb[e]; + // Exceptions of this tile with k in the segment: lanes load 32 at a time, each is broadcast; the owner lanes + // (row g or g + 8) add the exact product x[n][k] * w with x from shared memory. + { + int const g = lane >> 2, j = lane & 3; + int const kb = ku * 64, ke = (ku + seg) * 64; + uint32_t const* xs32 = reinterpret_cast(xs4); + 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(a.p.exc_pos + b0 + lane); + wv = __ldg(a.p.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); + // No early continue: every lane reaches each iteration's shuffles together. + 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; + uint32_t const word + = xs32[((((st >> 2) * 32 + n * 4 + ((rr & 7) >> 1)) * 4 + (st & 3)) * 2) + (rr >> 3)]; + float const xv = __uint_as_float((rr & 1) ? (word & 0xffff0000u) : (word << 16)); + float const d = __fmul_rn(xv, w0); + ca[q2] += hi ? 0.f : d; + ca[2 + q2] += hi ? d : 0.f; + } + } + __syncwarp(); + } + } + } + flush(a, warp, tile, ku * kUnit, (ku + seg) * kUnit, lane, ca); + done += seg; + ++tile; + ku = 0; + } +#endif +} + +} // namespace + +size_t targetHeadSmemBytes(int K) +{ + return static_cast(K / 64) * 64 * 16 + static_cast(kWarps) * kStages * 96 * 16; +} + +cudaError_t invokeTargetHead(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + if (p.M < 1 || p.M > kMaxRows || K % 64 != 0 || p.N < 1 || grid < 1) + return cudaErrorInvalidValue; + Args a; + a.p = p; + a.tiles = (p.N + 15) / 16; + a.ks = K / 16; + a.units = static_cast(a.tiles) * a.ks / kUnit; + // Every warp must own at least one unit: a split tile's finisher counts the warps between its first and last + // owner, so a warp with an empty range in between would leave the tile unstored. + if (a.units < kWarps) + return cudaErrorInvalidValue; + grid = static_cast(std::min(grid, a.units / kWarps)); + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(static_cast(grid)); + cfg.blockDim = dim3(kWarps * 32); + cfg.dynamicSmemBytes = targetHeadSmemBytes(K); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + auto kern = coded_kernel; + cudaError_t const err = cudaFuncSetAttribute( + kern, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(cfg.dynamicSmemBytes)); + if (err != cudaSuccess) + return err; + return cudaLaunchKernelEx(&cfg, kern, a); +} + +} // namespace targetHead +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/targetHead.h b/cpp/tensorrt_llm/kernels/minimaxM3/targetHead.h new file mode 100644 index 000000000000..84b68ee2152a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3/targetHead.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 +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3 +{ +//! The MiniMax-M3 target lm_head shard GEMV over a lossless 12-bit coding of its BF16 weight (verify rows <= 8, +//! sm_103a): BF16 x BF16 MMAs with FP32 accumulation over the full K, one BF16 rounding per logit. See targetHead.cu. +namespace targetHead +{ + +constexpr int kMaxRows = 8; // activation rows (verify rows) of one launch +constexpr int kWarps = 16; // warps per CTA +constexpr int kStages = 3; // coded units in flight per warp (cp.async ring) + +struct Params +{ + uint4 const* w; // [tiles][K / 64][3][32] coded weight (see targetHead.cu) + __nv_bfloat16 const* x; // [M, K] activation rows + long long x_stride; // elements between rows of x + __nv_bfloat16* out; // [M, N] logits + long long out_stride; // elements between rows of out + float* part; // [grid * kWarps][2][4 * 32] partial sums of tiles split between warps + int* counter; // [tiles] finishers per split tile, zero between calls (re-armed by the kernel) + int const* exc_off; // [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 exp_base; // exponent base << 7 in both BF16 halves (base << 7 | base << 23) + int M; // rows, 1..kMaxRows + int N; // logits per row (shard rows) +}; + +//! Shared memory a launch over hidden size K takes (the staged activation and the cp.async rings). +size_t targetHeadSmemBytes(int K); + +//! Launch over `grid` 512-thread CTAs (K a multiple of 64). +cudaError_t invokeTargetHead(Params const& p, int K, int grid, bool pdl, cudaStream_t stream); + +} // namespace targetHead +} // namespace minimaxM3 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/stepInputsScatter.cu b/cpp/tensorrt_llm/kernels/stepInputsScatter.cu new file mode 100644 index 000000000000..31616af5add0 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/stepInputsScatter.cu @@ -0,0 +1,69 @@ +/* + * 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/stepInputsScatter.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr int kThreads = 256; + +__device__ __forceinline__ void copyBytes(uint8_t* dst, uint8_t const* src, int64_t bytes) +{ + for (int64_t b = threadIdx.x; b < bytes; b += blockDim.x) + { + dst[b] = src[b]; + } +} + +// One CTA: the batches this serves are a step's few hundred bytes of metadata, where a single launch is the point and +// the copy itself is negligible. +__global__ void __launch_bounds__(kThreads) stepInputsScatterKernel(uint8_t const* __restrict__ stage) +{ + auto const* header = reinterpret_cast(stage); + int64_t const numHostCopies = header[0]; + int64_t const numDeviceCopies = header[1]; + int64_t const* desc = header + 2; + // Copies complete in descriptor order: one may overwrite or read another's destination. + for (int64_t i = 0; i < numHostCopies; ++i) + { + copyBytes(reinterpret_cast(desc[3 * i]), stage + desc[3 * i + 1], desc[3 * i + 2]); + __syncthreads(); + } + for (int64_t i = numHostCopies; i < numHostCopies + numDeviceCopies; ++i) + { + copyBytes(reinterpret_cast(desc[3 * i]), reinterpret_cast(desc[3 * i + 1]), + desc[3 * i + 2]); + __syncthreads(); + } +} + +} // namespace + +void invokeStepInputsScatter(uint8_t const* stage, cudaStream_t stream) +{ + stepInputsScatterKernel<<<1, kThreads, 0, stream>>>(stage); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/stepInputsScatter.h b/cpp/tensorrt_llm/kernels/stepInputsScatter.h new file mode 100644 index 000000000000..353197a479f9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/stepInputsScatter.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 +{ + +//! Write a packed batch of byte copies described by a staging buffer. +//! +//! The staging buffer starts with an int64 header `[numHostCopies, numDeviceCopies]`, followed by one +//! `(dst, src, bytes)` int64 triple per copy. A host copy's `src` is a byte offset into the staging buffer (its +//! bytes were uploaded with the header); a device copy's `src` is a device address. Every host copy is written +//! before any device copy, so a device copy may read a destination a host copy of the same batch filled. +void invokeStepInputsScatter(uint8_t const* stage, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index ab424318c862..3345690fa796 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -67,8 +67,24 @@ add_th_common_family( IndexerKCacheScatterOp.cpp IndexerTopKOp.cpp sparseKvCacheCompactOp.cpp + minimaxM3AttnC1Op.cpp + minimaxM3AttnMegaOp.cpp + minimaxM3DraftStepOp.cpp minimaxM3Fp8IndexerOp.cpp + minimaxM3FusedDecodeAttnOp.cpp + minimaxM3McAllreduceNormOp.cpp + minimaxM3MsaKvUpdateOp.cpp + minimaxM3Mxfp8LowMGemmOp.cpp + minimaxM3ProducerPwOp.cpp + minimaxM3RouterGemvOp.cpp + minimaxM3SampledAcceptOp.cpp minimaxM3SelectBlocksOp.cpp + minimaxM3SelectLongOp.cpp + minimaxM3ShardAcceptOp.cpp + minimaxM3SparseDecodeLcOp.cpp + minimaxM3SpecTailOp.cpp + minimaxM3TargetHeadOp.cpp + stepInputsScatterOp.cpp mlaRopeInplaceOp.cpp relativeAttentionBiasOp.cpp dsv3RopeOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3AttnC1Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3AttnC1Op.cpp new file mode 100644 index 000000000000..cfc915f939a9 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3AttnC1Op.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/minimaxM3/attnC1.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace dims = tensorrt_llm::kernels::minimax_m3_attn_c1; +using tensorrt_llm::kernels::MinimaxM3AttnC1Params; + +constexpr int64_t kHeadDim = 128; + +void checkByteTensor(at::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.element_size() == 1, name, " must be a CUDA tensor of one-byte elements"); +} + +} // namespace + +//! One launch for a one-request MiniMax-M3 MSA verify layer of 6 query tokens: the fused norm/RoPE/KV-insert + +//! identity-selection sparse decode + merge of minimax_m3_fused_decode_attn on the packed QKV rows, and the MXFP8 +//! o_proj GEMV, 96 CTAs that are each an attention CTA and an o_proj unit (fed through the MXFP8 hand-off). +//! +//! @param packed bf16 [6, 2560] QKV+index projection rows. +//! @param flags int32 [>= 32 + 96]: the per-CTA call counters (zero-initialized once); written. +//! @param mxfp8Data / mxfp8Scale the two call-parity buffers of the attention's MXFP8 hand-off (0xFF at rest). +//! @param y bf16 [T, 6144] o_proj output. arWs / arFlags / arRank / arWorld / arMaxTokens / arMode: the next +//! multicast all-reduce's exchange (device addresses) to push y into when arMode is 2; 0 for none. +void minimaxM3AttnC1(at::Tensor const& packed, at::Tensor& flags, at::Tensor& kvData, at::Tensor& kvScale, + at::Tensor& idxK, at::Tensor const& outCacheLoc, at::Tensor const& positionIds, at::Tensor const& blockTable, + at::Tensor const& seqLens, at::Tensor const& kvQuantScale, at::Tensor const& kGlobalScale, + at::Tensor const& vGlobalScale, at::Tensor const& qNormWeight, at::Tensor const& kNormWeight, + at::Tensor const& indexKNormWeight, at::Tensor const& rotaryCosSin, double eps, double smScale, + int64_t decodeQueryLen, int64_t tokensPerCluster, int64_t blocksPerCta, at::Tensor& output, at::Tensor& mxfp8Data, + at::Tensor& mxfp8Scale, bool usePdl, at::Tensor const& oWeight, at::Tensor const& oWeightScale, at::Tensor& y, + int64_t arWs, int64_t arFlags, int64_t arRank, int64_t arWorld, int64_t arMaxTokens, int64_t arMode) +{ + TORCH_CHECK(packed.is_cuda() && packed.scalar_type() == at::kBFloat16 && packed.dim() == 2 + && packed.size(1) == 20 * kHeadDim && packed.stride(1) == 1, + "packed must be CUDA bf16 [T, 2560] with contiguous rows"); + checkByteTensor(oWeight, "o_weight"); + checkByteTensor(oWeightScale, "o_weight_scale"); + checkByteTensor(kvData, "kv_data"); + checkByteTensor(kvScale, "kv_scale"); + checkByteTensor(idxK, "idx_k"); + checkByteTensor(mxfp8Data, "mxfp8_data"); + checkByteTensor(mxfp8Scale, "mxfp8_scale"); + TORCH_CHECK( + oWeight.dim() == 2 && oWeight.size(0) == dims::kOpN && oWeight.size(1) == dims::kOpK && oWeight.is_contiguous(), + "o_weight must be contiguous [6144, 2048]"); + TORCH_CHECK( + kvData.dim() == 5 && kvScale.dim() == 5 && idxK.dim() == 4, "kv_data, kv_scale and idx_k must be paged"); + for (auto const* t : {&outCacheLoc, &positionIds, &blockTable, &seqLens}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kInt, + "slot, position, block-table and length tensors " + "must be CUDA int32"); + } + TORCH_CHECK(blockTable.dim() == 2 && blockTable.stride(1) == 1, "block_table must be row-major 2-D"); + for (auto const* t : {&kvQuantScale, &kGlobalScale, &vGlobalScale, &rotaryCosSin}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kFloat, "scales and rotary_cos_sin must be CUDA fp32"); + } + for (auto const* t : {&qNormWeight, &kNormWeight, &indexKNormWeight}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kBFloat16 && t->numel() == kHeadDim, + "norm weights must be CUDA bf16 [128]"); + } + TORCH_CHECK(flags.is_cuda() && flags.scalar_type() == at::kInt && flags.numel() >= 32 + dims::kCtas, + "bad c1 flags: int32 [>= 128]"); + TORCH_CHECK(mxfp8Data.numel() == 2 * dims::kMaxTokens * dims::kOpK && mxfp8Scale.numel() == 2 * dims::kOpSfBytes + && mxfp8Data.is_contiguous() && mxfp8Scale.is_contiguous(), + "c1 MXFP8 hand-off"); + int64_t const numTokens = packed.size(0); + int64_t const numReqs = seqLens.size(0); + TORCH_CHECK(numTokens == numReqs * decodeQueryLen && numTokens <= dims::kMaxTokens, "token count"); + TORCH_CHECK(output.is_cuda() && output.scalar_type() == at::kBFloat16 && output.dim() == 2 + && output.size(0) >= numTokens && output.size(1) == 16 * kHeadDim && output.stride(1) == 1, + "output must be CUDA bf16 [T, 2048] with contiguous rows"); + TORCH_CHECK(y.is_cuda() && y.scalar_type() == at::kBFloat16 && y.dim() == 2 && y.size(0) == numTokens + && y.size(1) == dims::kOpN && y.is_contiguous(), + "y must be [M, 6144]"); + TORCH_CHECK(decodeQueryLen == 6 && tokensPerCluster == 1 && blocksPerCta == 1 && numReqs == 1, + "minimax_m3_attn_c1 is compiled for one request with six query tokens (geometry (6, 1, 1)), got ", numReqs, + " x (", decodeQueryLen, ", ", tokensPerCluster, ", ", blocksPerCta, ")"); + TORCH_CHECK(arMode == 0 || arMode == 2, "ar_mode must be 0 (none) or 2 (push)"); + + c10::cuda::CUDAGuard const deviceGuard{packed.device()}; + auto const* prop = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK( + prop->major == 10, "minimax_m3_attn_c1 requires an sm_100-family GPU, got sm_", prop->major, prop->minor); + + MinimaxM3AttnC1Params prm = {}; + prm.kv_data = reinterpret_cast(kvData.data_ptr()); + prm.d_page = kvData.stride(0); + prm.d_plane = kvData.stride(1); + prm.kv_scale = reinterpret_cast(kvScale.data_ptr()); + prm.s_page = kvScale.stride(0); + prm.s_plane = kvScale.stride(1); + prm.idx_k = reinterpret_cast(idxK.data_ptr()); + prm.i_page = idxK.stride(0); + prm.i_tok = idxK.stride(2); + prm.num_pages = kvData.size(0); + prm.out_cache_loc = outCacheLoc.data_ptr(); + prm.position_ids = positionIds.data_ptr(); + prm.block_table = blockTable.data_ptr(); + prm.bt_stride = blockTable.stride(0); + prm.bt_cols = static_cast(blockTable.size(1)); + prm.seq_lens = seqLens.data_ptr(); + prm.kv_quant_scale = kvQuantScale.data_ptr(); + prm.k_gs = kGlobalScale.data_ptr(); + prm.v_gs = vGlobalScale.data_ptr(); + prm.q_w = reinterpret_cast<__nv_bfloat16 const*>(qNormWeight.data_ptr()); + prm.k_w = reinterpret_cast<__nv_bfloat16 const*>(kNormWeight.data_ptr()); + prm.ik_w = reinterpret_cast<__nv_bfloat16 const*>(indexKNormWeight.data_ptr()); + prm.rope = rotaryCosSin.data_ptr(); + prm.eps = static_cast(eps); + prm.sm_scale = static_cast(smScale); + prm.out = reinterpret_cast<__nv_bfloat16*>(output.data_ptr()); + prm.out_stride = output.stride(0); + prm.q8 = reinterpret_cast(mxfp8Data.data_ptr()); + prm.sf = reinterpret_cast(mxfp8Scale.data_ptr()); + prm.num_sf_k_tiles = static_cast(output.size(1) / 128); + prm.packed = reinterpret_cast<__nv_bfloat16 const*>(packed.data_ptr()); + prm.packed_stride = packed.stride(0); + prm.num_tokens = static_cast(numTokens); + prm.flags = reinterpret_cast(flags.data_ptr()); + prm.wsf_o = reinterpret_cast(oWeightScale.data_ptr()); + prm.y_out = reinterpret_cast<__nv_bfloat16*>(y.data_ptr()); + prm.ar_ws = reinterpret_cast(arWs); + prm.ar_flags = reinterpret_cast(arFlags); + prm.ar_rank = static_cast(arRank); + prm.ar_world = static_cast(arWorld); + prm.ar_maxt = static_cast(arMaxTokens); + prm.ar_mode = static_cast(arMode); + + cudaError_t const err = tensorrt_llm::kernels::invokeMinimaxM3AttnC1(prm, oWeight.data_ptr(), + static_cast(decodeQueryLen), static_cast(tokensPerCluster), static_cast(blocksPerCta), + static_cast(numReqs), usePdl, at::cuda::getCurrentCUDAStream(packed.get_device())); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "minimax_m3_attn_c1 launch failed (qlen=", decodeQueryLen, " qg=", tokensPerCluster, + " bpc=", blocksPerCta, "): ", cudaGetErrorString(err)); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_attn_c1(Tensor packed, Tensor(a!) flags, Tensor(b!) kv_data, Tensor(c!) kv_scale, " + "Tensor(d!) idx_k, Tensor out_cache_loc, Tensor position_ids, Tensor block_table, Tensor seq_lens, " + "Tensor kv_quant_scale, Tensor k_global_scale, Tensor v_global_scale, Tensor q_norm_weight, " + "Tensor k_norm_weight, Tensor index_k_norm_weight, Tensor rotary_cos_sin, float eps, float sm_scale, " + "int decode_query_len, int tokens_per_cluster, int blocks_per_cta, Tensor(e!) output, Tensor(f!) mxfp8_data, " + "Tensor(g!) mxfp8_scale, bool use_pdl, Tensor o_weight, Tensor o_weight_scale, Tensor(h!) y, " + "int ar_workspace, int ar_flags, int ar_rank, int ar_world, int ar_max_tokens, int ar_mode) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_attn_c1", &tensorrt_llm::torch_ext::minimaxM3AttnC1); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3AttnMegaOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3AttnMegaOp.cpp new file mode 100644 index 000000000000..b94942e052f5 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3AttnMegaOp.cpp @@ -0,0 +1,207 @@ +/* + * 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/minimaxM3/attnMega.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace dims = tensorrt_llm::kernels::minimax_m3_attn_mega; +using tensorrt_llm::kernels::MinimaxM3AttnMegaParams; + +constexpr int64_t kHeadDim = 128; + +void checkByteTensor(at::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.element_size() == 1, name, " must be a CUDA tensor of one-byte elements"); +} + +} // namespace + +//! One launch for a one-request, one-token MiniMax-M3 MSA decode layer: the MXFP8 QKV GEMV of the layer input, the +//! fused norm/RoPE/KV-insert + identity-selection sparse decode + merge of minimax_m3_fused_decode_attn, and the MXFP8 +//! o_proj GEMV, joined by per-tile ready counters in `flags`. +//! +//! @param xq / xsf the layer input's E4M3 rows [T, 6144] and 128x4-swizzled UE8M0 scale atoms. +//! @param qkvWeight / qkvWeightScale E4M3 [2560, 6144] and its swizzled UE8M0 scales. +//! @param qkvOut bf16 [2, 8, 2560] call-parity buffers of the packed QKV rows (0xFFFF "not yet" at rest); written. +//! @param flags int32 [>= 32 + 112]: tile counters and the per-CTA call counters (zero-initialized once); written. +//! @param mxfp8Data / mxfp8Scale the two call-parity buffers of the attention's MXFP8 hand-off (0xFF at rest). +//! @param y bf16 [T, 6144] o_proj output. arWs / arFlags / arRank / arWorld / arMaxTokens / arMode: the next +//! multicast all-reduce's exchange (device addresses) to push y into when arMode is 2; 0 for none. +void minimaxM3AttnMega(at::Tensor const& xq, at::Tensor const& xsf, at::Tensor const& qkvWeight, + at::Tensor const& qkvWeightScale, at::Tensor& qkvOut, at::Tensor& flags, at::Tensor& kvData, at::Tensor& kvScale, + at::Tensor& idxK, at::Tensor const& outCacheLoc, at::Tensor const& positionIds, at::Tensor const& blockTable, + at::Tensor const& seqLens, at::Tensor const& kvQuantScale, at::Tensor const& kGlobalScale, + at::Tensor const& vGlobalScale, at::Tensor const& qNormWeight, at::Tensor const& kNormWeight, + at::Tensor const& indexKNormWeight, at::Tensor const& rotaryCosSin, double eps, double smScale, + int64_t decodeQueryLen, int64_t tokensPerCluster, int64_t blocksPerCta, at::Tensor& output, at::Tensor& mxfp8Data, + at::Tensor& mxfp8Scale, bool usePdl, at::Tensor const& oWeight, at::Tensor const& oWeightScale, at::Tensor& y, + int64_t arWs, int64_t arFlags, int64_t arRank, int64_t arWorld, int64_t arMaxTokens, int64_t arMode) +{ + checkByteTensor(xq, "xq"); + checkByteTensor(xsf, "xsf"); + checkByteTensor(qkvWeight, "qkv_weight"); + checkByteTensor(qkvWeightScale, "qkv_weight_scale"); + checkByteTensor(oWeight, "o_weight"); + checkByteTensor(oWeightScale, "o_weight_scale"); + checkByteTensor(kvData, "kv_data"); + checkByteTensor(kvScale, "kv_scale"); + checkByteTensor(idxK, "idx_k"); + checkByteTensor(mxfp8Data, "mxfp8_data"); + checkByteTensor(mxfp8Scale, "mxfp8_scale"); + TORCH_CHECK(xq.dim() == 2 && xq.size(1) == dims::kQkvK && xq.stride(1) == 1, "xq must be [T, 6144] row-major"); + TORCH_CHECK(qkvWeight.dim() == 2 && qkvWeight.size(0) == dims::kQkvN && qkvWeight.size(1) == dims::kQkvK + && qkvWeight.is_contiguous(), + "qkv_weight must be contiguous [2560, 6144]"); + TORCH_CHECK( + oWeight.dim() == 2 && oWeight.size(0) == dims::kOpN && oWeight.size(1) == dims::kOpK && oWeight.is_contiguous(), + "o_weight must be contiguous [6144, 2048]"); + TORCH_CHECK( + kvData.dim() == 5 && kvScale.dim() == 5 && idxK.dim() == 4, "kv_data, kv_scale and idx_k must be paged"); + for (auto const* t : {&outCacheLoc, &positionIds, &blockTable, &seqLens}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kInt, + "slot, position, block-table and length tensors " + "must be CUDA int32"); + } + TORCH_CHECK(blockTable.dim() == 2 && blockTable.stride(1) == 1, "block_table must be row-major 2-D"); + for (auto const* t : {&kvQuantScale, &kGlobalScale, &vGlobalScale, &rotaryCosSin}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kFloat, "scales and rotary_cos_sin must be CUDA fp32"); + } + for (auto const* t : {&qNormWeight, &kNormWeight, &indexKNormWeight}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kBFloat16 && t->numel() == kHeadDim, + "norm weights must be CUDA bf16 [128]"); + } + TORCH_CHECK(flags.is_cuda() && flags.scalar_type() == at::kInt && flags.numel() >= 32 + dims::kGrid, + "flags must be CUDA int32 [>= 144]"); + TORCH_CHECK(qkvOut.is_cuda() && qkvOut.element_size() == 2 && qkvOut.dim() == 3 && qkvOut.size(0) == 2 + && qkvOut.size(1) == dims::kMaxTokens && qkvOut.size(2) == dims::kQkvN && qkvOut.is_contiguous(), + "bad megakernel workspace: qkv_out must be contiguous [2, 8, 2560]"); + TORCH_CHECK(mxfp8Data.numel() == 2 * dims::kMaxTokens * dims::kOpK && mxfp8Scale.numel() == 2 * dims::kOpSfBytes + && mxfp8Data.is_contiguous() && mxfp8Scale.is_contiguous(), + "megakernel MXFP8 hand-off"); + int64_t const numTokens = xq.size(0); + int64_t const numReqs = seqLens.size(0); + TORCH_CHECK(numTokens == numReqs * decodeQueryLen && numTokens <= dims::kMaxTokens, "token count"); + TORCH_CHECK(output.is_cuda() && output.scalar_type() == at::kBFloat16 && output.dim() == 2 + && output.size(0) >= numTokens && output.size(1) == 16 * kHeadDim && output.stride(1) == 1, + "output must be CUDA bf16 [T, 2048] with contiguous rows"); + TORCH_CHECK(y.is_cuda() && y.scalar_type() == at::kBFloat16 && y.dim() == 2 && y.size(0) == numTokens + && y.size(1) == dims::kOpN && y.is_contiguous(), + "y must be [M, 6144]"); + TORCH_CHECK(decodeQueryLen == 1 && tokensPerCluster == 1 && blocksPerCta == 1 && numReqs == 1, + "minimax_m3_attn_mega is compiled for one request with one query token (geometry (1, 1, 1)), got ", numReqs, + " x (", decodeQueryLen, ", ", tokensPerCluster, ", ", blocksPerCta, ")"); + TORCH_CHECK(arMode == 0 || arMode == 2, "ar_mode must be 0 (none) or 2 (push)"); + + c10::cuda::CUDAGuard const deviceGuard{xq.device()}; + auto const* prop = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK( + prop->major == 10, "minimax_m3_attn_mega requires an sm_100-family GPU, got sm_", prop->major, prop->minor); + + MinimaxM3AttnMegaParams prm = {}; + prm.kv_data = reinterpret_cast(kvData.data_ptr()); + prm.d_page = kvData.stride(0); + prm.d_plane = kvData.stride(1); + prm.kv_scale = reinterpret_cast(kvScale.data_ptr()); + prm.s_page = kvScale.stride(0); + prm.s_plane = kvScale.stride(1); + prm.idx_k = reinterpret_cast(idxK.data_ptr()); + prm.i_page = idxK.stride(0); + prm.i_tok = idxK.stride(2); + prm.num_pages = kvData.size(0); + prm.out_cache_loc = outCacheLoc.data_ptr(); + prm.position_ids = positionIds.data_ptr(); + prm.block_table = blockTable.data_ptr(); + prm.bt_stride = blockTable.stride(0); + prm.bt_cols = static_cast(blockTable.size(1)); + prm.seq_lens = seqLens.data_ptr(); + prm.kv_quant_scale = kvQuantScale.data_ptr(); + prm.k_gs = kGlobalScale.data_ptr(); + prm.v_gs = vGlobalScale.data_ptr(); + prm.q_w = reinterpret_cast<__nv_bfloat16 const*>(qNormWeight.data_ptr()); + prm.k_w = reinterpret_cast<__nv_bfloat16 const*>(kNormWeight.data_ptr()); + prm.ik_w = reinterpret_cast<__nv_bfloat16 const*>(indexKNormWeight.data_ptr()); + prm.rope = rotaryCosSin.data_ptr(); + prm.eps = static_cast(eps); + prm.sm_scale = static_cast(smScale); + prm.out = reinterpret_cast<__nv_bfloat16*>(output.data_ptr()); + prm.out_stride = output.stride(0); + prm.q8 = reinterpret_cast(mxfp8Data.data_ptr()); + prm.sf = reinterpret_cast(mxfp8Scale.data_ptr()); + prm.num_sf_k_tiles = static_cast(output.size(1) / 128); + prm.xq = reinterpret_cast(xq.data_ptr()); + prm.ldxq = static_cast(xq.stride(0)); + prm.xsf_act = reinterpret_cast(xsf.data_ptr()); + prm.num_tokens = static_cast(numTokens); + prm.wsf = reinterpret_cast(qkvWeightScale.data_ptr()); + prm.qkv_out = reinterpret_cast<__nv_bfloat16*>(qkvOut.data_ptr()); + prm.flags = reinterpret_cast(flags.data_ptr()); + prm.wsf_o = reinterpret_cast(oWeightScale.data_ptr()); + prm.y_out = reinterpret_cast<__nv_bfloat16*>(y.data_ptr()); + prm.ar_ws = reinterpret_cast(arWs); + prm.ar_flags = reinterpret_cast(arFlags); + prm.ar_rank = static_cast(arRank); + prm.ar_world = static_cast(arWorld); + prm.ar_maxt = static_cast(arMaxTokens); + prm.ar_mode = static_cast(arMode); + + cudaError_t const err + = tensorrt_llm::kernels::invokeMinimaxM3AttnMega(prm, qkvWeight.data_ptr(), oWeight.data_ptr(), + static_cast(decodeQueryLen), static_cast(tokensPerCluster), static_cast(blocksPerCta), + static_cast(numReqs), usePdl, at::cuda::getCurrentCUDAStream(xq.get_device())); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "minimax_m3_attn_mega launch failed (qlen=", decodeQueryLen, " qg=", tokensPerCluster, + " bpc=", blocksPerCta, "): ", cudaGetErrorString(err)); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_attn_mega(Tensor xq, Tensor xsf, Tensor qkv_weight, Tensor qkv_weight_scale, Tensor(a!) qkv_out, " + "Tensor(b!) flags, Tensor(c!) kv_data, Tensor(d!) kv_scale, Tensor(e!) idx_k, Tensor out_cache_loc, " + "Tensor position_ids, Tensor block_table, Tensor seq_lens, Tensor kv_quant_scale, Tensor k_global_scale, " + "Tensor v_global_scale, Tensor q_norm_weight, Tensor k_norm_weight, Tensor index_k_norm_weight, " + "Tensor rotary_cos_sin, float eps, float sm_scale, int decode_query_len, int tokens_per_cluster, " + "int blocks_per_cta, Tensor(f!) output, Tensor(g!) mxfp8_data, Tensor(h!) mxfp8_scale, bool use_pdl, " + "Tensor o_weight, Tensor o_weight_scale, Tensor(i!) y, int ar_workspace, int ar_flags, int ar_rank, " + "int ar_world, int ar_max_tokens, int ar_mode) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_attn_mega", &tensorrt_llm::torch_ext::minimaxM3AttnMega); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3DraftStepOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3DraftStepOp.cpp new file mode 100644 index 000000000000..6bfc88d79213 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3DraftStepOp.cpp @@ -0,0 +1,198 @@ +/* + * 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/minimaxM3/draftStep.h" + +#include +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace ds = tensorrt_llm::kernels::minimaxM3::draftStep; + +using OptTensor = std::optional; + +void checkArch(torch::Tensor const& t) +{ + TORCH_CHECK(t.is_cuda(), "minimax_m3_draft_step expects CUDA tensors"); + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_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 minimaxM3DraftStep(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) +{ + 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))}; + 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_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) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_draft_step", &tensorrt_llm::torch_ext::minimaxM3DraftStep); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3FusedDecodeAttnOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3FusedDecodeAttnOp.cpp new file mode 100644 index 000000000000..e431a62f0028 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3FusedDecodeAttnOp.cpp @@ -0,0 +1,175 @@ +/* + * 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/minimaxM3/fusedDecodeAttn.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +using tensorrt_llm::kernels::MinimaxM3FusedDecodeAttnParams; + +constexpr int64_t kNumHeads = 16; +constexpr int64_t kHeadDim = 128; +constexpr int64_t kPackedWidth = (kNumHeads + 4) * kHeadDim; + +void checkByteTensor(at::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.element_size() == 1, name, " must be a CUDA tensor of one-byte elements"); +} + +} // namespace + +//! Fused per-head Gemma RMSNorm + partial NeoX RoPE + NVFP4 K/V and FP8 index-K cache insert + identity-selection +//! sparse decode attention over the paged NVFP4 cache + split-K merge, for the MiniMax-M3 MSA sparse layers in +//! short-sequence decode (every request attends at most 16 KV blocks of 128 tokens). +//! +//! @param packed bf16 [T, 2560] QKV+index projection rows ([Q x16 | K | V | idxQ | idxK] x 128), T = R * qlen. +//! @param kvData uint8 [pages, 2, 1, 128, 64] HND NVFP4 pool (K plane 0, V plane 1); the new tokens are written. +//! @param kvScale uint8 [pages, 2, 1, 128, 8] E4M3 block scales (K linear, V 4x4 token quads); written. +//! @param idxK one-byte [pages, 1, 128, 128] FP8 index-K cache; written. +//! @param output bf16 [T, 2048] attention output; mxfp8Data / mxfp8Scale its MXFP8 copy (E4M3 rows and 128x4-swizzled +//! UE8M0 scales) for the o_proj GEMM. +void minimaxM3FusedDecodeAttn(at::Tensor const& packed, at::Tensor& kvData, at::Tensor& kvScale, at::Tensor& idxK, + at::Tensor const& outCacheLoc, at::Tensor const& positionIds, at::Tensor const& blockTable, + at::Tensor const& seqLens, at::Tensor const& kvQuantScale, at::Tensor const& kGlobalScale, + at::Tensor const& vGlobalScale, at::Tensor const& qNormWeight, at::Tensor const& kNormWeight, + at::Tensor const& indexKNormWeight, at::Tensor const& rotaryCosSin, double eps, double smScale, + int64_t decodeQueryLen, int64_t tokensPerCluster, int64_t blocksPerCta, at::Tensor& output, at::Tensor& mxfp8Data, + at::Tensor& mxfp8Scale, bool usePdl) +{ + TORCH_CHECK(packed.is_cuda() && packed.scalar_type() == at::kBFloat16 && packed.dim() == 2 + && packed.size(1) == kPackedWidth && packed.stride(1) == 1, + "packed must be CUDA bf16 [T, 2560] with contiguous rows"); + checkByteTensor(kvData, "kv_data"); + checkByteTensor(kvScale, "kv_scale"); + checkByteTensor(idxK, "idx_k"); + TORCH_CHECK( + kvData.dim() == 5 && kvScale.dim() == 5 && idxK.dim() == 4, "kv_data, kv_scale and idx_k must be paged"); + for (auto const* t : {&outCacheLoc, &positionIds, &blockTable, &seqLens}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kInt, + "slot, position, block-table and length tensors " + "must be CUDA int32"); + } + TORCH_CHECK(blockTable.dim() == 2 && blockTable.stride(1) == 1, "block_table must be row-major 2-D"); + for (auto const* t : {&kvQuantScale, &kGlobalScale, &vGlobalScale, &rotaryCosSin}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kFloat, "scales and rotary_cos_sin must be CUDA fp32"); + } + for (auto const* t : {&qNormWeight, &kNormWeight, &indexKNormWeight}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kBFloat16 && t->numel() == kHeadDim, + "norm weights must be CUDA bf16 [128]"); + } + TORCH_CHECK(output.is_cuda() && output.scalar_type() == at::kBFloat16 && output.dim() == 2 + && output.size(1) == kNumHeads * kHeadDim && output.stride(1) == 1, + "output must be CUDA bf16 [T, 2048] with contiguous rows"); + checkByteTensor(mxfp8Data, "mxfp8_data"); + checkByteTensor(mxfp8Scale, "mxfp8_scale"); + int64_t const numTokens = packed.size(0); + int64_t const numReqs = seqLens.size(0); + TORCH_CHECK(numTokens == numReqs * decodeQueryLen, "packed rows must equal num_requests * decode_query_len"); + TORCH_CHECK(output.size(0) >= numTokens && mxfp8Data.dim() == 2 && mxfp8Data.size(0) >= numTokens + && mxfp8Data.size(1) == kNumHeads * kHeadDim && mxfp8Data.is_contiguous() && mxfp8Scale.is_contiguous(), + "mxfp8_data must be contiguous [>= T, 2048]"); + TORCH_CHECK(tensorrt_llm::kernels::minimaxM3FusedDecodeAttnSupports(static_cast(decodeQueryLen), + static_cast(tokensPerCluster), static_cast(blocksPerCta)), + "minimax_m3_fused_decode_attn: geometry (", decodeQueryLen, ", ", tokensPerCluster, ", ", blocksPerCta, + ") is not compiled"); + if (numTokens == 0) + { + return; + } + + c10::cuda::CUDAGuard const deviceGuard{packed.device()}; + auto const* prop = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(prop->major == 10, "minimax_m3_fused_decode_attn requires an sm_100-family GPU, got sm_", prop->major, + prop->minor); + + MinimaxM3FusedDecodeAttnParams prm = {}; + prm.packed = reinterpret_cast<__nv_bfloat16 const*>(packed.data_ptr()); + prm.packed_stride = packed.stride(0); + prm.kv_data = reinterpret_cast(kvData.data_ptr()); + prm.d_page = kvData.stride(0); + prm.d_plane = kvData.stride(1); + prm.kv_scale = reinterpret_cast(kvScale.data_ptr()); + prm.s_page = kvScale.stride(0); + prm.s_plane = kvScale.stride(1); + prm.idx_k = reinterpret_cast(idxK.data_ptr()); + prm.i_page = idxK.stride(0); + prm.i_tok = idxK.stride(2); + prm.num_pages = kvData.size(0); + prm.out_cache_loc = outCacheLoc.data_ptr(); + prm.position_ids = positionIds.data_ptr(); + prm.block_table = blockTable.data_ptr(); + prm.bt_stride = blockTable.stride(0); + prm.bt_cols = static_cast(blockTable.size(1)); + prm.seq_lens = seqLens.data_ptr(); + prm.kv_quant_scale = kvQuantScale.data_ptr(); + prm.k_gs = kGlobalScale.data_ptr(); + prm.v_gs = vGlobalScale.data_ptr(); + prm.q_w = reinterpret_cast<__nv_bfloat16 const*>(qNormWeight.data_ptr()); + prm.k_w = reinterpret_cast<__nv_bfloat16 const*>(kNormWeight.data_ptr()); + prm.ik_w = reinterpret_cast<__nv_bfloat16 const*>(indexKNormWeight.data_ptr()); + prm.rope = rotaryCosSin.data_ptr(); + prm.eps = static_cast(eps); + prm.sm_scale = static_cast(smScale); + prm.out = reinterpret_cast<__nv_bfloat16*>(output.data_ptr()); + prm.out_stride = output.stride(0); + prm.q8 = reinterpret_cast(mxfp8Data.data_ptr()); + prm.sf = reinterpret_cast(mxfp8Scale.data_ptr()); + prm.num_sf_k_tiles = static_cast(output.size(1) / 128); + + cudaError_t const err = tensorrt_llm::kernels::invokeMinimaxM3FusedDecodeAttn(prm, static_cast(decodeQueryLen), + static_cast(tokensPerCluster), static_cast(blocksPerCta), static_cast(numReqs), usePdl, + at::cuda::getCurrentCUDAStream(packed.get_device())); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "minimax_m3_fused_decode_attn launch failed (qlen=", decodeQueryLen, + " qg=", tokensPerCluster, " bpc=", blocksPerCta, "): ", cudaGetErrorString(err)); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_fused_decode_attn(Tensor packed, Tensor(a!) kv_data, Tensor(b!) kv_scale, Tensor(c!) idx_k, " + "Tensor out_cache_loc, Tensor position_ids, Tensor block_table, Tensor seq_lens, Tensor kv_quant_scale, " + "Tensor k_global_scale, Tensor v_global_scale, Tensor q_norm_weight, Tensor k_norm_weight, " + "Tensor index_k_norm_weight, Tensor rotary_cos_sin, float eps, float sm_scale, int decode_query_len, " + "int tokens_per_cluster, int blocks_per_cta, Tensor(d!) output, Tensor(e!) mxfp8_data, " + "Tensor(f!) mxfp8_scale, bool use_pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_fused_decode_attn", &tensorrt_llm::torch_ext::minimaxM3FusedDecodeAttn); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3McAllreduceNormOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3McAllreduceNormOp.cpp new file mode 100644 index 000000000000..3f4558f8a330 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3McAllreduceNormOp.cpp @@ -0,0 +1,159 @@ +/* + * 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/minimaxM3/mcAllreduceNorm.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 multicast all-reduce + 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. +//! @return [normed, updated residual] or, with `quant`, [normed, updated residual, E4M3 rows, scale bytes]. +std::vector minimaxM3McAllreduceNorm(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) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sm == 100 || sm == 103, "minimax_m3_mc_allreduce_norm requires SM100 or SM103, got SM", sm); + TORCH_CHECK(cluster == kCluster, "minimax_m3_mc_allreduce_norm: only cluster 8 is built, got ", cluster); + TORCH_CHECK(tensorrt_llm::kernels::isMinimaxM3McAllreduceNormSupported( + static_cast(maxt), static_cast(opts)), + "minimax_m3_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::MinimaxM3McAllreduceNormParams 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); + tensorrt_llm::kernels::invokeMinimaxM3McAllreduceNorm(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_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[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_mc_allreduce_norm", &tensorrt_llm::torch_ext::minimaxM3McAllreduceNorm); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3MsaKvUpdateOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3MsaKvUpdateOp.cpp new file mode 100644 index 000000000000..dfea9d3af430 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3MsaKvUpdateOp.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/minimaxM3/msaKvUpdate.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +void checkInt32Cuda(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda(), "minimax_m3_msa_kv_update: ", name, " must be a CUDA tensor"); + TORCH_CHECK( + t.scalar_type() == torch::kInt32, "minimax_m3_msa_kv_update: ", name, " must be int32, got ", t.scalar_type()); +} + +} // namespace + +//! MiniMax-M3 MSA KV-length correction after the overlap scheduler moved kv_lens (see invokeMinimaxM3MsaKvUpdate). +//! +//! Every tensor is int32 on one CUDA device. block_table is [rows, ncols] with unit column stride; every other +//! operand is contiguous with at least `batch` (per request) or `total_q` (per query token) elements, and cu_kv_lens +//! at least `batch + 1`. Writes msa_seq_lens, out_cache_loc and, when given, cu_kv_lens and n_valid_blocks. +void minimaxM3MsaKvUpdate(torch::Tensor const& kvLens, torch::Tensor const& kvLensStaged, torch::Tensor const& seqLens, + torch::Tensor const& qBatchRow, torch::Tensor const& qIntra, torch::Tensor const& blockTable, + int64_t tokensPerBlock, int64_t batch, int64_t totalQ, torch::Tensor& msaSeqLens, + std::optional cuKvLens, torch::Tensor& outCacheLoc, std::optional nValidBlocks) +{ + checkInt32Cuda(kvLens, "kv_lens"); + checkInt32Cuda(kvLensStaged, "kv_lens_staged"); + checkInt32Cuda(seqLens, "seq_lens"); + checkInt32Cuda(qBatchRow, "q_batch_row"); + checkInt32Cuda(qIntra, "q_intra"); + checkInt32Cuda(blockTable, "block_table"); + checkInt32Cuda(msaSeqLens, "msa_seq_lens"); + checkInt32Cuda(outCacheLoc, "out_cache_loc"); + TORCH_CHECK(blockTable.dim() == 2 && blockTable.stride(1) == 1, + "minimax_m3_msa_kv_update: block_table must be 2-D with unit column stride"); + TORCH_CHECK(tokensPerBlock > 0, "minimax_m3_msa_kv_update: tokens_per_block must be positive"); + TORCH_CHECK(batch >= 0 && totalQ >= 0, "minimax_m3_msa_kv_update: batch and total_q must be non-negative"); + TORCH_CHECK(kvLens.numel() >= batch && kvLensStaged.numel() >= batch && seqLens.numel() >= batch + && msaSeqLens.numel() >= batch, + "minimax_m3_msa_kv_update: a per-request operand has fewer than batch elements"); + TORCH_CHECK(qBatchRow.numel() >= totalQ && qIntra.numel() >= totalQ && outCacheLoc.numel() >= totalQ, + "minimax_m3_msa_kv_update: a per-token operand has fewer than total_q elements"); + int32_t* cuKvPtr = nullptr; + if (cuKvLens.has_value()) + { + checkInt32Cuda(*cuKvLens, "cu_kv_lens"); + TORCH_CHECK(cuKvLens->numel() >= batch + 1, "minimax_m3_msa_kv_update: cu_kv_lens needs batch + 1 elements"); + cuKvPtr = cuKvLens->data_ptr(); + } + int32_t* nValidPtr = nullptr; + if (nValidBlocks.has_value()) + { + checkInt32Cuda(*nValidBlocks, "n_valid_blocks"); + TORCH_CHECK(nValidBlocks->numel() >= totalQ, "minimax_m3_msa_kv_update: n_valid_blocks needs total_q elements"); + nValidPtr = nValidBlocks->data_ptr(); + } + + c10::cuda::CUDAGuard const deviceGuard{kvLens.device()}; + auto const stream = at::cuda::getCurrentCUDAStream(kvLens.get_device()); + cudaError_t const err = tensorrt_llm::kernels::invokeMinimaxM3MsaKvUpdate(kvLens.data_ptr(), + kvLensStaged.data_ptr(), seqLens.data_ptr(), qBatchRow.data_ptr(), + qIntra.data_ptr(), blockTable.data_ptr(), blockTable.stride(0), + static_cast(blockTable.size(1)), static_cast(tokensPerBlock), static_cast(batch), + static_cast(totalQ), msaSeqLens.data_ptr(), cuKvPtr, outCacheLoc.data_ptr(), + nValidPtr, stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_msa_kv_update launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_msa_kv_update(Tensor kv_lens, Tensor kv_lens_staged, Tensor seq_lens, Tensor q_batch_row, " + "Tensor q_intra, Tensor block_table, int tokens_per_block, int batch, int total_q, Tensor(a!) msa_seq_lens, " + "Tensor(b!)? cu_kv_lens, Tensor(c!) out_cache_loc, Tensor(d!)? n_valid_blocks) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_msa_kv_update", &tensorrt_llm::torch_ext::minimaxM3MsaKvUpdate); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3Mxfp8LowMGemmOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3Mxfp8LowMGemmOp.cpp new file mode 100644 index 000000000000..c8f6b4a0355d --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3Mxfp8LowMGemmOp.cpp @@ -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. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3/mxfp8LowMGemm.h" + +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! MiniMax-M3 fused MXFP8 activation quantization + low-M MXFP8 GEMM (decode). +//! +//! @param x BF16 [M, K] activation, unit column stride, 16-byte aligned rows (row stride a multiple of 8). +//! @param w E4M3 [N, K] weight, contiguous. +//! @param wsf UE8M0 weight scales in the 128x4 swizzled layout of `mxfp8_quantize(w, True)`. +//! @param warps, clusterK, rowTiles The kernel configuration; see `isMinimaxM3Mxfp8LowMGemmSupported`. +//! @return BF16 [M, N]. +torch::Tensor minimaxM3Mxfp8LowMGemm(torch::Tensor const& x, torch::Tensor const& w, torch::Tensor const& wsf, + int64_t warps, int64_t clusterK, int64_t rowTiles, bool pdl) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sm == 100 || sm == 103, "minimax_m3_mxfp8_low_m_gemm requires SM100 or SM103, got SM", sm); + 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"); + TORCH_CHECK(w.scalar_type() == at::kFloat8_e4m3fn && w.is_contiguous() && w.dim() == 2 && w.size(1) == x.size(1), + "w must be a contiguous e4m3 [N, K] tensor"); + TORCH_CHECK(wsf.scalar_type() == at::kByte, "wsf must be uint8 UE8M0 scales"); + TORCH_CHECK( + reinterpret_cast(x.data_ptr()) % 16 == 0 && x.stride(0) % 8 == 0, "x rows must be 16-byte aligned"); + TORCH_CHECK( + reinterpret_cast(w.data_ptr()) % 16 == 0 && reinterpret_cast(wsf.data_ptr()) % 16 == 0, + "w and wsf must be 16-byte aligned"); + TORCH_CHECK(tensorrt_llm::kernels::isMinimaxM3Mxfp8LowMGemmSupported(static_cast(warps), + static_cast(clusterK), static_cast(rowTiles), static_cast(x.size(0))), + "minimax_m3_mxfp8_low_m_gemm: unsupported configuration (warps ", warps, ", cluster_k ", clusterK, + ", row_tiles ", rowTiles, ") at M=", x.size(0)); + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto out = torch::empty({x.size(0), w.size(0)}, x.options()); + tensorrt_llm::kernels::invokeMinimaxM3Mxfp8LowMGemm(reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()), + static_cast(x.stride(0)), static_cast(x.size(0)), static_cast(w.size(0)), + static_cast(x.size(1)), reinterpret_cast(w.data_ptr()), + reinterpret_cast(wsf.data_ptr()), reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), + static_cast(out.stride(0)), static_cast(warps), static_cast(clusterK), + static_cast(rowTiles), pdl, at::cuda::getCurrentCUDAStream(x.get_device())); + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_mxfp8_low_m_gemm(Tensor x, Tensor w, Tensor wsf, int warps, int cluster_k, int row_tiles, " + "bool pdl) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_mxfp8_low_m_gemm", &tensorrt_llm::torch_ext::minimaxM3Mxfp8LowMGemm); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3ProducerPwOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3ProducerPwOp.cpp new file mode 100644 index 000000000000..ec73b498e40e --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3ProducerPwOp.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/minimaxM3/producerPw.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace pw = tensorrt_llm::kernels::minimaxM3::producerPw; + +constexpr int64_t kHeadDim = 128; +constexpr int64_t kPageSize = 128; + +void checkCuda(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda(), "producer_pw: ", name, " must be a CUDA tensor"); +} + +} // namespace + +//! The NVFP4 QKV + indexer norm/RoPE/KV-insert producer (same arguments and outputs as +//! minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert for 1 KV head, 1 index head, head_dim 128, rotary_dim 64): +//! returns (q FP8 [T, headsQ, 128], index_q FP8 [T, 1, 128]) and writes the NVFP4 K/V and FP8 index-K of each token's +//! cache slot. `early` triggers the dependents on entry. +std::vector minimaxM3ProducerPw(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, + int64_t headsQ, bool pdl, bool early) +{ + 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(), + "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, + "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_producer_pw requires SM100 or newer"); + TORCH_CHECK(headsQ == 8 || headsQ == 16, "producer_pw: 8 or 16 query heads"); + int64_t const totalHeads = headsQ + 4; + TORCH_CHECK(packed.scalar_type() == at::kBFloat16 && packed.dim() == 2 && packed.size(1) == totalHeads * kHeadDim + && packed.is_contiguous(), + "producer_pw: packed must be contiguous BF16 [T, (headsQ + 4) * 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) == 1 && kvData.size(3) == kPageSize && kvData.size(4) == kHeadDim / 2, + "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) == 1 && kvScale.size(3) == kPageSize + && kvScale.size(4) == kHeadDim / 16, + "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), + "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), + "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), + "producer_pw: pool planes or pages overlap"); + TORCH_CHECK(kvData.stride(0) % 2 == 0 && kvData.stride(1) % 2 == 0, + "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, + "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, + "producer_pw: out_cache_loc must be int32 [>= T]"); + TORCH_CHECK( + positionIds.scalar_type() == at::kInt && positionIds.is_contiguous() && positionIds.numel() >= numTokens, + "producer_pw: position_ids must be int32 [>= T]"); + TORCH_CHECK(kvQuantScale.scalar_type() == at::kFloat && kvQuantScale.is_contiguous() && kvQuantScale.numel() >= 3, + "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, + "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, + "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, 1, 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::invokeProducerPw(p, static_cast(numTokens), static_cast(headsQ), pdl, early, + at::cuda::getCurrentCUDAStream(packed.get_device())); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "producer_pw launch: ", cudaGetErrorString(err)); + } + return {qOut, indexQOut}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_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, int heads_q, " + "bool pdl, bool early) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_producer_pw", &tensorrt_llm::torch_ext::minimaxM3ProducerPw); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3RouterGemvOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3RouterGemvOp.cpp new file mode 100644 index 000000000000..059d29e43d85 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3RouterGemvOp.cpp @@ -0,0 +1,149 @@ +/* + * 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/minimaxM3/routerGemv.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kGemvMaxM = 16; +constexpr int64_t kTcMaxM = 48; +constexpr int64_t kTcExperts = 128; + +void checkSmVersion(char const* op) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sm == 100 || sm == 103, op, " requires SM100 or SM103, 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, + "x must be a 2-D bf16 CUDA tensor with unit column stride"); + TORCH_CHECK(x.size(0) >= 1 && x.size(0) <= maxM, "x must hold 1..", maxM, " rows, got ", x.size(0)); + TORCH_CHECK(w.scalar_type() == at::kFloat && w.is_contiguous() && w.dim() == 2 && w.size(1) == x.size(1), + "w must be a contiguous fp32 [E, K] tensor"); + TORCH_CHECK(x.size(1) % 1024 == 0 && (x.stride(0) * 2) % 16 == 0, "K must be a multiple of 1024 with 16-byte rows"); +} + +} // namespace + +//! fp32 router logits `x @ w^T` for 1..16 bf16 rows (exact fp32 FMA, one CTA per expert). +torch::Tensor minimaxM3RouterGemv(torch::Tensor const& x, torch::Tensor const& w, int64_t rows, bool pdl) +{ + checkSmVersion("minimax_m3_router_gemv"); + checkActivation(x, w, kGemvMaxM); + TORCH_CHECK(rows == 1, "minimax_m3_router_gemv: only rows=1 is built, got ", rows); + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto out = torch::empty({x.size(0), w.size(0)}, x.options().dtype(at::kFloat)); + tensorrt_llm::kernels::invokeMinimaxM3RouterGemv(reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()), + static_cast(x.stride(0)), static_cast(x.size(0)), static_cast(x.size(1)), + w.data_ptr(), static_cast(w.size(0)), out.data_ptr(), nullptr, nullptr, nullptr, pdl, + at::cuda::getCurrentCUDAStream(x.get_device())); + return out; +} + +//! Router logits as `minimax_m3_router_gemv` plus the NVFP4 activation: [logits fp32 [M, E], packed e2m1 +//! [M, K / 2] uint8, linear E4M3 scales [M * K / 16] uint8]. +std::vector minimaxM3RouterGemvFp4( + torch::Tensor const& x, torch::Tensor const& w, torch::Tensor const& sfScale, bool pdl) +{ + checkSmVersion("minimax_m3_router_gemv_fp4"); + checkActivation(x, w, kGemvMaxM); + TORCH_CHECK(sfScale.scalar_type() == at::kFloat && sfScale.numel() == 1 && sfScale.is_cuda(), + "sf_scale must be a one-element fp32 CUDA tensor"); + int64_t const M = x.size(0), K = x.size(1); + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto out = torch::empty({M, w.size(0)}, x.options().dtype(at::kFloat)); + auto q = torch::empty({M, K / 2}, x.options().dtype(at::kByte)); + auto sf = torch::empty({M * (K / 16)}, x.options().dtype(at::kByte)); + tensorrt_llm::kernels::invokeMinimaxM3RouterGemv(reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()), + static_cast(x.stride(0)), static_cast(M), static_cast(K), w.data_ptr(), + static_cast(w.size(0)), out.data_ptr(), q.data_ptr(), sf.data_ptr(), + sfScale.data_ptr(), pdl, at::cuda::getCurrentCUDAStream(x.get_device())); + return {out, q, sf}; +} + +//! TF32 tensor-core router logits for 128 experts and 1..48 rows: [logits fp32 [M, 128]] or, with `sf_scale`, +//! [logits, packed e2m1 [M, K / 2], linear E4M3 scales [M * K / 16]]. +std::vector minimaxM3RouterTc( + torch::Tensor const& x, torch::Tensor const& w, std::optional const& sfScale, bool pdl) +{ + checkSmVersion("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); + TORCH_CHECK(M >= 1 && M <= kTcMaxM && w.scalar_type() == at::kFloat && w.dim() == 2 && w.size(0) == kTcExperts + && w.size(1) == K && w.is_contiguous(), + "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()}; + auto out = torch::empty({M, w.size(0)}, x.options().dtype(at::kFloat)); + torch::Tensor q, sf; + uint8_t* qp = nullptr; + uint8_t* sfp = nullptr; + float const* scp = nullptr; + if (sfScale.has_value()) + { + TORCH_CHECK(sfScale->scalar_type() == at::kFloat && sfScale->numel() == 1 && sfScale->is_cuda(), + "sf_scale must be a one-element fp32 CUDA tensor"); + q = torch::empty({M, K / 2}, x.options().dtype(at::kByte)); + sf = torch::empty({M * (K / 16)}, x.options().dtype(at::kByte)); + qp = q.data_ptr(); + sfp = sf.data_ptr(); + scp = sfScale->data_ptr(); + } + tensorrt_llm::kernels::invokeMinimaxM3RouterTc(reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()), + static_cast(x.stride(0)), static_cast(M), static_cast(K), w.data_ptr(), + out.data_ptr(), qp, sfp, scp, pdl, at::cuda::getCurrentCUDAStream(x.get_device())); + if (sfScale.has_value()) + { + return {out, q, sf}; + } + return {out}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +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[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_router_gemv", &tensorrt_llm::torch_ext::minimaxM3RouterGemv); + m.impl("minimax_m3_router_gemv_fp4", &tensorrt_llm::torch_ext::minimaxM3RouterGemvFp4); + m.impl("minimax_m3_router_tc", &tensorrt_llm::torch_ext::minimaxM3RouterTc); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3SampledAcceptOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3SampledAcceptOp.cpp new file mode 100644 index 000000000000..1c2183647374 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3SampledAcceptOp.cpp @@ -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. + */ + +#include "tensorrt_llm/kernels/minimaxM3/sampledAccept.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace sa = tensorrt_llm::kernels::minimaxM3::sampledAccept; + +//! uint64 words of the sampled-accept exchange workspace for `world` ranks. +int64_t minimaxM3SampledAcceptWsWords(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 minimaxM3SampledAccept(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) +{ + 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; + 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_sampled_accept_ws_words(int world) -> int", + &tensorrt_llm::torch_ext::minimaxM3SampledAcceptWsWords); + m.def( + "minimax_m3_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) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sampled_accept", &tensorrt_llm::torch_ext::minimaxM3SampledAccept); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3SelectLongOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3SelectLongOp.cpp new file mode 100644 index 000000000000..2def696ce875 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3SelectLongOp.cpp @@ -0,0 +1,80 @@ +/* + * 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/minimaxM3/selectLong.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. +void minimaxM3SelectLong(torch::Tensor const& scores, torch::Tensor const& nValidBlocks, torch::Tensor& out, + int64_t initBlocks, int64_t localBlocks, bool usePdl, bool earlyTrigger) +{ + TORCH_CHECK(scores.is_cuda() && scores.scalar_type() == torch::kFloat32 && scores.dim() == 3, + "minimax_m3_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_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::kMinimaxM3SelectLongMaxBlocks, + "minimax_m3_select_long supports at most ", tensorrt_llm::kernels::kMinimaxM3SelectLongMaxBlocks, + " 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_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::invokeMinimaxM3SelectLong(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, stream); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "minimax_m3_select_long launch failed: ", cudaGetErrorString(err)); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_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) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_select_long", &tensorrt_llm::torch_ext::minimaxM3SelectLong); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3ShardAcceptOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3ShardAcceptOp.cpp new file mode 100644 index 000000000000..5b4ca3e6004f --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3ShardAcceptOp.cpp @@ -0,0 +1,118 @@ +/* + * 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/minimaxM3/shardAccept.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace sa = tensorrt_llm::kernels::minimaxM3::shardAccept; + +using OptTensor = std::optional; + +void checkArch(torch::Tensor const& t) +{ + TORCH_CHECK(t.is_cuda(), "minimax_m3_shard_accept expects CUDA tensors"); + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_shard_accept requires SM100 or newer"); +} + +void check(cudaError_t err) +{ + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, cudaGetErrorString(err)); + } +} + +cudaStream_t currentStream(torch::Tensor const& t) +{ + return at::cuda::getCurrentCUDAStream(t.get_device()); +} + +//! Per-request accepted counts for fractional forced acceptance (int32 [nreq]), or nullptr. +int const* forceVecPtr(OptTensor const& forceVec, int nreq) +{ + if (!forceVec.has_value()) + { + return nullptr; + } + TORCH_CHECK(forceVec->scalar_type() == at::kInt && forceVec->is_contiguous() && forceVec->numel() >= nreq, + "accept: force_vec must be contiguous int32 [nreq]"); + return forceVec->data_ptr(); +} + +} // namespace + +//! Greedy verification of nreq * (draft_len + 1) <= 48 verify rows from this rank's BF16 logits shard [rows, V / tp] +//! (global ids vocabOffset..), exchanged across the TP ranks through the multicast slots wsMcPtr / wsUc, plus the +//! strict acceptance (accepted tokens, counts, first-pass gather rows, optional nextNew[:, 0]). +void minimaxM3ShardAccept(torch::Tensor const& logits, int64_t vocabOffset, torch::Tensor& keys, + torch::Tensor& counters, torch::Tensor& wsUc, int64_t wsMcPtr, 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, OptTensor const& nextNew, int64_t numCtas, OptTensor const& forceVec) +{ + checkArch(logits); + c10::cuda::CUDAGuard const deviceGuard{logits.device()}; + TORCH_CHECK(logits.scalar_type() == at::kBFloat16 && logits.dim() == 2 && logits.stride(1) == 1); + int const M = logits.size(0); + int const nreq = accTokens.size(0); + int const k1 = accTokens.size(1); + TORCH_CHECK(M >= 1 && M <= sa::kMaxRows && nreq * k1 == M && nreq <= sa::kThreads); + TORCH_CHECK(draft.scalar_type() == at::kInt && batchIdx.scalar_type() == at::kInt && accTokens.is_contiguous()); + sa::AcceptArgs acc{draft.data_ptr(), batchIdx.data_ptr(), accTokens.data_ptr(), + numAcc.data_ptr(), reinterpret_cast(gatherIds.data_ptr()), k1 - 1, nreq, + static_cast(forceTotal), nextNew.has_value() ? nextNew->data_ptr() : nullptr, + nextNew.has_value() ? static_cast(nextNew->stride(0)) : 0, forceVecPtr(forceVec, nreq)}; + check(sa::invokeShardAccept(reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()), M, + static_cast(logits.size(1)), static_cast(logits.stride(0)), static_cast(vocabOffset), + reinterpret_cast(keys.data_ptr()), reinterpret_cast(counters.data_ptr()), + reinterpret_cast(wsUc.data_ptr()), reinterpret_cast(wsMcPtr), + static_cast(rank), static_cast(world), acc, static_cast(numCtas), currentStream(logits))); +} + +} // 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_shard_accept(Tensor logits, int vocab_offset, Tensor(a!) keys, Tensor(b!) counters, " + "Tensor(c!) ws_uc, int ws_mc_ptr, 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, " + "int num_ctas, Tensor? force_vec) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_shard_accept", &tensorrt_llm::torch_ext::minimaxM3ShardAccept); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3SparseDecodeLcOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3SparseDecodeLcOp.cpp new file mode 100644 index 000000000000..62a56779f88b --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3SparseDecodeLcOp.cpp @@ -0,0 +1,139 @@ +/* + * 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/minimaxM3/sparseDecodeLc.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_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 minimaxM3SparseDecodeLc(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) +{ + TORCH_CHECK(earlyTrigger, "minimax_m3_sparse_decode_lc is instantiated with early_trigger=True only"); + TORCH_CHECK(q.is_cuda(), "minimax_m3_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_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_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_sparse_decode_lc: global scales must be float32"); + TORCH_CHECK(topk.scalar_type() == torch::kInt32 && topk.dim() == 2 && topk.size(1) == kTopK, + "minimax_m3_sparse_decode_lc: topk must be int32 [T, 16]"); + TORCH_CHECK(blockTable.scalar_type() == torch::kInt32 && seqLens.scalar_type() == torch::kInt32, + "minimax_m3_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_sparse_decode_lc: out must be bfloat16 [T, 2048] with unit last-dim stride"); + TORCH_CHECK(q8.has_value() == sf.has_value(), "minimax_m3_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::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); + + auto const stream = at::cuda::getCurrentCUDAStream(q.get_device()); + cudaError_t const err = tensorrt_llm::kernels::invokeMinimaxM3SparseDecodeLc( + prm, static_cast(q.size(0)), linearSoftmax, mxfp8, earlyTrigger, usePdl, stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sparse_decode_lc launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_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) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sparse_decode_lc", &tensorrt_llm::torch_ext::minimaxM3SparseDecodeLc); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3SpecTailOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3SpecTailOp.cpp new file mode 100644 index 000000000000..cea04a45cd86 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3SpecTailOp.cpp @@ -0,0 +1,202 @@ +/* + * 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/minimaxM3/specTail.h" + +#include +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace st = tensorrt_llm::kernels::minimaxM3::specTail; + +void check(cudaError_t err) +{ + TORCH_CHECK(err == cudaSuccess, cudaGetErrorString(err)); +} + +cudaStream_t currentStream(torch::Tensor const& t) +{ + return at::cuda::getCurrentCUDAStream(t.get_device()); +} + +__nv_bfloat16 const* bf16In(torch::Tensor const& t) +{ + return reinterpret_cast<__nv_bfloat16 const*>(t.data_ptr()); +} + +__nv_bfloat16* bf16Out(torch::Tensor const& t) +{ + return reinterpret_cast<__nv_bfloat16*>(t.data_ptr()); +} + +//! Register-tile count of the norm kernels for hidden size H (H / 8 elements over kThreads threads). +int normIters(int H) +{ + return (H / 8 + st::kThreads - 1) / st::kThreads; +} + +st::IntView intView(torch::Tensor const& t) +{ + return st::IntView{t.data_ptr(), t.stride(0), t.stride(1), static_cast(t.size(1))}; +} + +} // namespace + +//! out = [rmsnorm(embed(ids)) * wEmbed, rmsnorm(hidden[rows]) * wHidden] [M, 2H]; residual = hidden[rows]. +void minimaxM3DraftInput(torch::Tensor const& ids, torch::Tensor const& table, torch::Tensor const& hidden, + std::optional const& rows, torch::Tensor const& wEmbed, torch::Tensor const& wHidden, + double epsEmbed, double epsHidden, torch::Tensor& out, std::optional const& residual, bool pdl) +{ + c10::cuda::CUDAGuard const deviceGuard{table.device()}; + int const M = ids.size(0); + int const H = table.size(1); + TORCH_CHECK(ids.scalar_type() == at::kInt && ids.is_contiguous() && table.scalar_type() == at::kBFloat16 + && table.is_contiguous() && H % 8 == 0); + TORCH_CHECK( + hidden.scalar_type() == at::kBFloat16 && hidden.dim() == 2 && hidden.is_contiguous() && hidden.size(1) == H); + TORCH_CHECK(out.scalar_type() == at::kBFloat16 && out.size(0) == M && out.size(1) == 2 * H && out.is_contiguous()); + TORCH_CHECK(wEmbed.scalar_type() == at::kBFloat16 && wEmbed.numel() == H && wHidden.scalar_type() == at::kBFloat16 + && wHidden.numel() == H); + TORCH_CHECK(!rows.has_value() || (rows->scalar_type() == at::kLong && rows->is_contiguous() && rows->size(0) == M)); + TORCH_CHECK(rows.has_value() || hidden.size(0) >= M, "draft_input: hidden must hold a row per id"); + TORCH_CHECK(!residual.has_value() + || (residual->scalar_type() == at::kBFloat16 && residual->is_contiguous() && residual->size(0) == M + && residual->size(1) == H), + "draft_input: residual must be contiguous [M, H]"); + TORCH_CHECK(normIters(H) <= st::kMaxIters, "draft_input: hidden size too large"); + check(st::invokeDraftInput(ids.data_ptr(), static_cast(table.size(0)), bf16In(table), bf16In(hidden), + rows.has_value() ? reinterpret_cast(rows->data_ptr()) : nullptr, bf16In(wEmbed), + bf16In(wHidden), static_cast(epsEmbed), static_cast(epsHidden), H, M, bf16Out(out), + residual.has_value() ? bf16Out(*residual) : nullptr, pdl, currentStream(table))); +} + +//! Per-chunk RMSNorm of src [M, >= C*H] (row stride may exceed C*H) with weights w [C, H] into dst [M, C*H]. +void minimaxM3ChunkNorm(torch::Tensor const& src, torch::Tensor const& w, double eps, torch::Tensor& dst, bool pdl) +{ + c10::cuda::CUDAGuard const deviceGuard{src.device()}; + int const M = src.size(0); + int const C = w.size(0); + int const H = w.size(1); + TORCH_CHECK(src.scalar_type() == at::kBFloat16 && w.scalar_type() == at::kBFloat16 && src.stride(1) == 1); + TORCH_CHECK(src.size(1) >= C * H && dst.is_contiguous() && dst.numel() == static_cast(M) * C * H); + TORCH_CHECK(H % 8 == 0 && src.stride(0) % 8 == 0); + TORCH_CHECK(normIters(H) <= st::kMaxIters, "chunk_norm: hidden size too large"); + check(st::invokeChunkNorm(bf16In(src), src.stride(0), bf16In(w), static_cast(eps), H, C, M, bf16Out(dst), + pdl, currentStream(src))); +} + +//! dst [M, H] (row stride may exceed H) = rmsnorm(bf16(h + r)) * w, r optional. +void minimaxM3CaptureNorm(torch::Tensor const& h, std::optional const& r, torch::Tensor const& w, + double eps, torch::Tensor& dst) +{ + c10::cuda::CUDAGuard const deviceGuard{h.device()}; + int const M = h.size(0); + int const H = w.size(0); + TORCH_CHECK( + h.scalar_type() == at::kBFloat16 && w.scalar_type() == at::kBFloat16 && dst.scalar_type() == at::kBFloat16); + TORCH_CHECK(h.dim() == 2 && h.size(1) == H && h.stride(1) == 1 && dst.dim() == 2 && dst.size(0) == M + && dst.size(1) == H && dst.stride(1) == 1 && H % 8 == 0 && h.stride(0) % 8 == 0 && dst.stride(0) % 8 == 0); + TORCH_CHECK(!r.has_value() + || (r->sizes() == h.sizes() && r->stride(1) == 1 && r->stride(0) % 8 == 0 + && r->scalar_type() == at::kBFloat16)); + TORCH_CHECK(normIters(H) <= st::kMaxIters, "capture_norm: hidden size too large"); + check(st::invokeCaptureNorm(bf16In(h), h.stride(0), r.has_value() ? bf16In(*r) : nullptr, + r.has_value() ? r->stride(0) : 0, bf16In(w), static_cast(eps), H, M, bf16Out(dst), dst.stride(0), + currentStream(h))); +} + +//! The overlap scheduler's step inputs from the previous step's device outputs (see step_gather_kernel). +void minimaxM3StepGather(torch::Tensor const& newTok, torch::Tensor const& nextDraft, torch::Tensor const& lens, + torch::Tensor const& slots, torch::Tensor const& ppi, int64_t P, int64_t T, int64_t D, torch::Tensor& ids, + int64_t ids0, torch::Tensor& draft, int64_t d0, torch::Tensor& posOff, torch::Tensor& kvOff, int64_t b0) +{ + c10::cuda::CUDAGuard const deviceGuard{newTok.device()}; + for (auto const* t : std::initializer_list{ + &newTok, &nextDraft, &lens, &slots, &ppi, &ids, &draft, &posOff, &kvOff}) + { + TORCH_CHECK(t->scalar_type() == at::kInt, "step_gather: int32 tensors expected"); + } + TORCH_CHECK(newTok.dim() >= 2 && nextDraft.dim() == 2 && lens.is_contiguous() && posOff.is_contiguous() + && kvOff.is_contiguous() && ids.is_contiguous() && draft.is_contiguous()); + int const work = std::max(std::max(P * T, P * D), std::max(posOff.numel(), kvOff.numel())); + check(st::invokeStepGather(newTok.data_ptr(), newTok.stride(0), newTok.stride(1), nextDraft.data_ptr(), + nextDraft.stride(0), nextDraft.stride(1), lens.data_ptr(), slots.data_ptr(), ppi.data_ptr(), + static_cast(P), static_cast(T), static_cast(D), ids.data_ptr(), static_cast(ids0), + draft.data_ptr(), static_cast(d0), posOff.data_ptr(), static_cast(posOff.numel()), + kvOff.data_ptr(), static_cast(kvOff.numel()), static_cast(b0), work, currentStream(newTok))); +} + +//! The spec sampler's slot stores from one step's outputs (see slot_scatter_kernel). +void minimaxM3SlotScatter(torch::Tensor const& slots, torch::Tensor const& sa, torch::Tensor& da, + torch::Tensor const& sb, torch::Tensor& db, torch::Tensor const& sl, torch::Tensor& dl, torch::Tensor const& sd, + torch::Tensor& dd) +{ + c10::cuda::CUDAGuard const deviceGuard{slots.device()}; + for (auto const* t : std::initializer_list{&sa, &da, &sb, &db, &sl, &dl, &sd, &dd}) + { + TORCH_CHECK(t->scalar_type() == at::kInt && t->dim() <= 2, "slot_scatter: int32 tensors expected"); + } + TORCH_CHECK(slots.scalar_type() == at::kLong); + int const n = slots.numel(); + if (n == 0) + { + return; + } + check(st::invokeSlotScatter(reinterpret_cast(slots.data_ptr()), n, intView(sa), intView(da), + intView(sb), intView(db), sl.data_ptr(), dl.data_ptr(), intView(sd), intView(dd), + currentStream(slots))); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_draft_input(Tensor ids, Tensor table, Tensor hidden, Tensor? rows, Tensor w_embed, " + "Tensor w_hidden, float eps_embed, float eps_hidden, Tensor(a!) out, Tensor(b!)? residual, bool pdl) -> ()"); + m.def("minimax_m3_chunk_norm(Tensor src, Tensor w, float eps, Tensor(a!) dst, bool pdl) -> ()"); + m.def("minimax_m3_capture_norm(Tensor h, Tensor? r, Tensor w, float eps, Tensor(a!) dst) -> ()"); + m.def( + "minimax_m3_step_gather(Tensor new_tok, Tensor next_draft, Tensor lens, Tensor slots, Tensor ppi, int P, " + "int T, int D, Tensor(a!) ids, int ids0, Tensor(b!) draft, int d0, Tensor(c!) pos_off, Tensor(d!) kv_off, " + "int b0) -> ()"); + m.def( + "minimax_m3_slot_scatter(Tensor slots, Tensor sa, Tensor(a!) da, Tensor sb, Tensor(b!) db, Tensor sl, " + "Tensor(c!) dl, Tensor sd, Tensor(d!) dd) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_draft_input", &tensorrt_llm::torch_ext::minimaxM3DraftInput); + m.impl("minimax_m3_chunk_norm", &tensorrt_llm::torch_ext::minimaxM3ChunkNorm); + m.impl("minimax_m3_capture_norm", &tensorrt_llm::torch_ext::minimaxM3CaptureNorm); + m.impl("minimax_m3_step_gather", &tensorrt_llm::torch_ext::minimaxM3StepGather); + m.impl("minimax_m3_slot_scatter", &tensorrt_llm::torch_ext::minimaxM3SlotScatter); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3TargetHeadOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3TargetHeadOp.cpp new file mode 100644 index 000000000000..448989e08c3a --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3TargetHeadOp.cpp @@ -0,0 +1,123 @@ +/* + * 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/minimaxM3/targetHead.h" + +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace th = tensorrt_llm::kernels::minimaxM3::targetHead; + +void checkArch(torch::Tensor const& t) +{ + TORCH_CHECK(t.is_cuda(), "minimax_m3_target_head expects CUDA tensors"); + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_target_head requires SM100 or newer"); +} + +} // namespace + +//! out [M, n] (BF16) = x [M <= 8, K] (BF16) x W^T for the 12-bit coded BF16 shard W [n, K]: w the coded stream +//! (uint8 [tiles * K / 64 * 1536]), exc_off int32 [tiles + 1] / exc_pos int32 [E] / exc_w fp32 [E, 2] its exception +//! lists, exp_base its exponent window's base. part (fp32 [grid * 16 * 256]) and counter (int32 [tiles], zero) are +//! the split-tile scratch. +void minimaxM3TargetHead(torch::Tensor const& w, torch::Tensor const& x, torch::Tensor& out, torch::Tensor& part, + torch::Tensor& counter, torch::Tensor const& excOff, torch::Tensor const& excPos, torch::Tensor const& excW, + int64_t expBase, int64_t n, int64_t grid, bool pdl) +{ + checkArch(x); + for (torch::Tensor const* t : + std::initializer_list{&w, &out, &part, &counter, &excOff, &excPos, &excW}) + { + TORCH_CHECK(t->device() == x.device(), "target_head: 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, "target_head: x must be BF16 [M, K]"); + int const M = static_cast(x.size(0)); + int const K = static_cast(x.size(1)); + TORCH_CHECK(M >= 1 && M <= th::kMaxRows, "target_head: 1..", th::kMaxRows, " rows"); + TORCH_CHECK(K % 64 == 0 && K / 16 <= 0xFFFF, "target_head: K must be a multiple of 64"); + TORCH_CHECK(x.stride(1) == 1 && x.stride(0) % 8 == 0 && reinterpret_cast(x.data_ptr()) % 16 == 0, + "target_head: x rows must be unit-stride and 16-byte aligned"); + TORCH_CHECK(n >= 16 && n % 16 == 0, "target_head: n must be a positive multiple of 16 (whole tiles)"); + int64_t const tiles = (n + 15) / 16; + TORCH_CHECK(w.scalar_type() == at::kByte && w.is_contiguous() && w.numel() == tiles * 16 * K * 3 / 2, + "target_head: w must be the contiguous uint8 coded stream of [ceil(n / 16) * 16, K]"); + TORCH_CHECK(out.scalar_type() == at::kBFloat16 && out.dim() == 2 && out.size(0) >= M && out.size(1) >= n + && out.stride(1) == 1, + "target_head: out must be BF16 [M, >= n] with unit column stride"); + TORCH_CHECK(grid >= 1, "target_head: grid must be positive"); + TORCH_CHECK(part.scalar_type() == at::kFloat && part.is_contiguous() && part.numel() >= grid * th::kWarps * 2 * 128, + "target_head: part must be fp32 [grid * 16 * 256]"); + TORCH_CHECK(counter.scalar_type() == at::kInt && counter.is_contiguous() && counter.numel() >= tiles, + "target_head: counter must be int32 [tiles]"); + TORCH_CHECK(excOff.scalar_type() == at::kInt && excOff.is_contiguous() && excOff.numel() == tiles + 1, + "target_head: exc_off must be int32 [tiles + 1]"); + TORCH_CHECK(excPos.scalar_type() == at::kInt && excPos.is_contiguous(), "target_head: exc_pos must be int32"); + TORCH_CHECK(excW.scalar_type() == at::kFloat && excW.is_contiguous() && excW.numel() == 2 * excPos.numel(), + "target_head: exc_w must be fp32 [E, 2]"); + TORCH_CHECK(expBase >= 0 && expBase <= 240, "target_head: exponent base out of range"); + th::Params p{}; + p.w = reinterpret_cast(w.data_ptr()); + p.x = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.x_stride = x.stride(0); + p.out = reinterpret_cast<__nv_bfloat16*>(out.data_ptr()); + p.out_stride = out.stride(0); + p.part = part.data_ptr(); + p.counter = counter.data_ptr(); + p.exc_off = excOff.data_ptr(); + p.exc_pos = reinterpret_cast(excPos.data_ptr()); + p.exc_w = reinterpret_cast(excW.data_ptr()); + p.exp_base = static_cast((expBase << 7) | (expBase << 23)); + p.M = M; + p.N = static_cast(n); + cudaError_t const err + = th::invokeTargetHead(p, K, static_cast(grid), pdl, at::cuda::getCurrentCUDAStream(x.get_device())); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, cudaGetErrorString(err)); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_target_head(Tensor w, Tensor x, Tensor(a!) out, Tensor(b!) part, Tensor(c!) counter, " + "Tensor exc_off, Tensor exc_pos, Tensor exc_w, int exp_base, int n, int grid, bool pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_target_head", &tensorrt_llm::torch_ext::minimaxM3TargetHead); +} diff --git a/cpp/tensorrt_llm/thop/stepInputsScatterOp.cpp b/cpp/tensorrt_llm/thop/stepInputsScatterOp.cpp new file mode 100644 index 000000000000..e8d168d68849 --- /dev/null +++ b/cpp/tensorrt_llm/thop/stepInputsScatterOp.cpp @@ -0,0 +1,57 @@ +/* + * 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/stepInputsScatter.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! Issue a batch of packed host-to-device and device-to-device copies in one launch. +//! +//! @param stage Contiguous uint8 CUDA tensor laid out as described in kernels/stepInputsScatter.h. Its destinations are +//! raw device addresses: the caller guarantees each one is live and that no other work on the stream writes it between +//! staging and this call. +void stepInputsScatter(torch::Tensor const& stage) +{ + TORCH_CHECK(stage.is_cuda(), "step_inputs_scatter expects a CUDA staging buffer"); + TORCH_CHECK(stage.scalar_type() == torch::kUInt8, "step_inputs_scatter expects a uint8 staging buffer"); + TORCH_CHECK(stage.is_contiguous(), "step_inputs_scatter expects a contiguous staging buffer"); + TORCH_CHECK(stage.numel() >= 16, "step_inputs_scatter staging buffer is shorter than its header"); + c10::cuda::CUDAGuard const deviceGuard{stage.device()}; + tensorrt_llm::kernels::invokeStepInputsScatter( + stage.data_ptr(), at::cuda::getCurrentCUDAStream(stage.get_device())); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("step_inputs_scatter(Tensor stage) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("step_inputs_scatter", &tensorrt_llm::torch_ext::stepInputsScatter); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py index c5140e7c912b..7dc4b711fd91 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py @@ -32,6 +32,7 @@ ModelingV2Mode, NoTarget, modeling_v2_resolve, + resolved_mode, step_contract_enabled, ) @@ -41,5 +42,6 @@ "ModelingV2Mode", "NoTarget", "modeling_v2_resolve", + "resolved_mode", "step_contract_enabled", ] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index cdb288f7220f..66a63a7bf5b2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -37,7 +37,7 @@ import enum import os from importlib import import_module -from typing import TYPE_CHECKING, Optional, Tuple +from typing import TYPE_CHECKING, Dict, Optional, Tuple from tensorrt_llm.logger import logger @@ -51,6 +51,7 @@ MODELING_V2_ROUTERS = { "GptOssForCausalLM": "models.gpt_oss.routing", "DeepseekV3ForCausalLM": "models.deepseek_v3.routing", + "MiniMaxM3SparseForConditionalGeneration": "models.minimax_m3.routing", } @@ -101,6 +102,20 @@ class NoTarget(Exception): """ +# The mode each target class was last resolved under (the loader builds the class right after resolving it). +_RESOLVED_MODES: Dict[type, ModelingV2Mode] = {} + + +def resolved_mode(target: type) -> ModelingV2Mode: + """The mode ``modeling_v2_resolve`` returned ``target`` under; ``auto`` for a class built without it. + + A target reads it at construction: under ``require`` its in-target dispatch + raises where it would otherwise fall back to a built-in path inside the + shapes it claims. + """ + return _RESOLVED_MODES.get(target, ModelingV2Mode.AUTO) + + def routing_module(arch: str): """Import and return the routing module for ``arch``, or None.""" name = MODELING_V2_ROUTERS.get(arch) @@ -143,7 +158,7 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti raise NoTarget( f"no routing module for this architecture; routed architectures: {known}" ) - return routing.route( + target = routing.route( pretrained_config, config.mapping, _device_compute_capability(), config.spec_config ) except NoTarget as why: @@ -153,3 +168,6 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti f"modeling_v2: no target fits {arch!r} ({why}); using the built-in implementation" ) return None + if target is not None: + _RESOLVED_MODES[target] = mode + return target diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/__init__.py new file mode 100644 index 000000000000..a3dfe4a6ce39 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/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 +"""MiniMaxM3SparseForConditionalGeneration targets.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/__init__.py new file mode 100644 index 000000000000..fdb324a42eda --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/__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 NVFP4 / sm_103 / tp4.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/dispatch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/dispatch.py new file mode 100644 index 000000000000..47b00f57275d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/dispatch.py @@ -0,0 +1,580 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-step fast paths of this target: the rows each one takes, the catalog entry that certifies them, and +what happens when one cannot run. + +One target serves every concurrency with the same code; which fast path a step takes is a pure function of its +shape, written here once (``plan_for``, ``moe_path``, ``router_path``, ``attention_path``, ``linear_path``) and +called by the forward. Every row range is a named constant below, and ``FAST_PATHS`` pairs each with the catalog +entry whose ``CELLS`` certify it: ``test_modeling_v2_m3_dispatch.py`` checks that every row count (or shape key) a +range covers is one of the entry's cells. + +Two kinds of not taking a fast path, kept apart: + +* **uncovered**: the step is outside every covered range -- prefill, mixed steps, generation wider than the ranges, + an eager decode step (the GEMV dispatch is a decode-graph one), a long-context request, an Eagle3 capture layer + that runs the generic layer by design. ``uncovered`` logs the reason once per (path, reason); the caller runs the + generic path in every mode. +* **missed**: the step is inside a covered range but a precondition is missing -- a workspace, a push descriptor or + a compiled engine not built, or a kernel declining operands it is certified for. Under ``modeling_v2='require'`` + ``missed`` raises ``FastPathUnavailable`` naming the path, the rows and the precondition; otherwise it logs once + per (path, reason) and the caller falls back. + +Compiled kernels (the MoE engines, the decode GEMVs with their push variants, the index scorer) are built by +``build_targets`` 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, kernel modules) 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``) rather +than a silent capture of the fallback. +""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Callable, Dict, FrozenSet, Hashable, Iterable, Mapping, Optional, Tuple + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper +from tensorrt_llm.logger import logger + +from .ops.attention.m3_attn_c1 import m3_attn_c1 +from .ops.attention.m3_attn_mega import m3_attn_mega +from .ops.attention.m3_dense_decode import m3_dense_decode +from .ops.attention.m3_fused_decode_attn import m3_fused_decode_attn +from .ops.attention.m3_index_score import m3_index_score +from .ops.attention.m3_producer import m3_producer +from .ops.attention.m3_select_long import MAX_BLOCKS as SELECT_LONG_MAX_BLOCKS +from .ops.attention.m3_select_long import m3_select_long +from .ops.attention.m3_sparse_decode import m3_sparse_decode +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.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_target_head import MAX_ROWS as TARGET_HEAD_KERNEL_MAX_ROWS +from .ops.gemm.m3_target_head import m3_target_head +from .ops.gemm.mxfp8_low_m_gemm import mxfp8_low_m_gemm +from .ops.moe.m3_moe16_engine import m3_moe16_engine +from .ops.moe.m3_moe_engine import m3_moe_engine +from .ops.moe.m3_router_gemv import m3_router_gemv +from .ops.moe.m3_router_gemv_fp4 import m3_router_gemv_fp4 +from .ops.moe.m3_router_tc import m3_router_tc +from .ops.norm.m3_embedding_firstnorm import m3_embedding_firstnorm +from .ops.spec.m3_draft_input import m3_draft_input +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_shard_accept import m3_shard_accept + +# -------------------------------------------------------------------------------------------------------------------- +# Row ranges +# -------------------------------------------------------------------------------------------------------------------- + +# The largest generation step the decode target takes: 8 requests verifying 1 + 3 drafts and their padding rows. +DECODE_MAX_TOKENS = 48 + + +@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 (kernels/m3_mc_allreduce.BUCKETS). +DECODE_PLANS = ( + DecodePlan("tiny", 8, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), + DecodePlan("small", 16, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), + DecodePlan("verify", DECODE_MAX_TOKENS, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), +) + +# 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. +DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({1: 16384}) + +# 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 + +# The fused embedding + first RMSNorm kernel: one owner per row count and index dtype. +EMBED_FIRSTNORM_MAX_TOKENS = 8 +# 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 = 16 +AR_PUSH_MODE = 2 +# The 1..8-token MoE engine (router, routed experts and the shared expert in one launch) pushes its sum into the MoE +# boundary's all-reduce (prepushed mode 2). +MOE_ENGINE_MAX_TOKENS = 8 +MOE_PUSH_MAX_TOKENS = MOE_ENGINE_MAX_TOKENS +# The 9..13-row engine (router, routed and shared experts in one kernel); its bf16 sum goes to the boundary. +MOE16_MIN_TOKENS = 9 +MOE16_MAX_TOKENS = 13 +# 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). +MC_FINALIZE_MIN_TOKENS = 9 +# The router: the fp32 GEMV at one row, the tensor-core router (fused with the routed experts' NVFP4 activation +# quantization) at 2..48 rows. GB300 tp4, behind a PDL producer, with the NVFP4 activation: 2.90 / 2.93 / 3.20 / +# 3.82 / 7.36 us at 1 / 4 / 8 / 12 / 48 tokens against the fp32 GEMV's 2.85 / 3.96 / 4.88 / 6.01 / - us. +ROUTER_GEMV_MAX_TOKENS = 1 +ROUTER_TC_MAX_TOKENS = 48 +# The fused norm/RoPE/KV-insert/sparse-decode/merge kernel: up to 8 query rows per request and 16 query rows per +# step (above, its 16-CTA clusters run in waves and the three-kernel path is faster: GB300 -13% fAL spec c8). +FUSED_ATTN_MAX_QUERY_LEN = 8 +FUSED_ATTN_MAX_TOKENS = 16 +# One request verifying a draft of 6 query rows runs attention + o_proj as one launch (96 CTAs: 6 clusters of 16). +ATTN_C1_QUERY_LEN = 6 +# 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 (kernels/m3_shard_accept), at the configured draft length. +SHARD_ACCEPT_DRAFT_LEN = 3 +SHARD_ACCEPT_MAX_REQUESTS = DECODE_MAX_TOKENS // (SHARD_ACCEPT_DRAFT_LEN + 1) +# 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 (kernels/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. +TARGET_HEAD_MAX_ROWS = TARGET_HEAD_KERNEL_MAX_ROWS +# 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 = 8 +# 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 (kernels/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 (kernels/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 = True +DRAFT_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN +DRAFT_STEP_MAX_REQUESTS = DRAFT_STEP_KERNEL_MAX_TOKENS +# The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows +# (kernels/triton_sparse_decode.MQ_MIN_TOKENS). +SPARSE_DECODE_MQ_MIN_TOKENS = 17 +# 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 +LONG_SELECT_MAX_BLOCKS = SELECT_LONG_MAX_BLOCKS +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; on wider steps the early-launched attention CTAs crowd +# the select's CTAs off the SMs (8 requests x 4 rows: +3.1 us per layer), so they keep the trigger at the select's exit. +SELECT_EARLY_MAX_ROWS = 8 +# 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 (larger batches +# are bandwidth bound, where the extra early tiles only crowd the prefetch). +INDEX_SCORE_EARLY_MAX_REQUESTS = 1 + +# 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. +_ALL_ROWS = frozenset(range(1, DECODE_MAX_TOKENS + 1)) +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, +} +# 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, +} +# The fused low-M MXFP8 GEMM (projections marked for it), per projection: its built configurations' row limit +# (kernels/mxfp8_low_m_gemm). The bf16 GEMV serves every one of these rows first; the GEMM takes them only if it +# declines. +LOW_M_GEMM_MAX_TOKENS: Dict[Tuple[int, int], int] = { + QKV_SPARSE: 8, + O_PROJ: 8, + DENSE_GATE_UP: 1, + DENSE_DOWN: 8, +} + +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 _fused_attn_keys() -> FrozenSet[Tuple[int, int]]: + return frozenset( + (qlen, reqs) + for qlen in range(1, FUSED_ATTN_MAX_QUERY_LEN + 1) + for reqs in range(1, FUSED_ATTN_MAX_TOKENS // qlen + 1) + ) + + +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"), +) +# One request decoding one token on an MXFP8 layer input: QKV GEMV + attention + o_proj in one launch. +ATTN_MEGA = RowRange("attn_mega", frozenset({(1, 1)}), m3_attn_mega, _const((1, 1))) +ATTN_C1 = RowRange( + "attn_c1", frozenset({(ATTN_C1_QUERY_LEN, 1)}), m3_attn_c1, _const((ATTN_C1_QUERY_LEN, 1)) +) +FUSED_DECODE_ATTN = RowRange( + "fused_decode_attn", + _fused_attn_keys(), + m3_fused_decode_attn, + lambda spec: ((spec["qlen"], len(spec["seq_lens"])),), +) +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",), +) +LONG_SELECT = RowRange( + "long_select", frozenset(LONG_SELECT_KEY_WIDTHS), m3_select_long, _long_select_widths +) +MC_ALLREDUCE = RowRange( + "mc_allreduce_norm", + _rows(1, DECODE_MAX_TOKENS), + 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), +) +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, DECODE_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"), +) +ROUTER_GEMV = RowRange( + "m3_router_gemv", _rows(1, ROUTER_GEMV_MAX_TOKENS), m3_router_gemv, _spec_rows("m") +) +ROUTER_GEMV_FP4 = RowRange( + "m3_router_gemv_fp4", _rows(1, ROUTER_GEMV_MAX_TOKENS), m3_router_gemv_fp4, _spec_rows("m") +) +ROUTER_TC = RowRange( + "m3_router_tc", + _rows(ROUTER_GEMV_MAX_TOKENS + 1, ROUTER_TC_MAX_TOKENS), + m3_router_tc, + _spec_rows("m"), +) +# 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"]),), +) +LOW_M_GEMM = RowRange( + "mxfp8_low_m_gemm", + frozenset((nk, m) for nk, hi in LOW_M_GEMM_MAX_TOKENS.items() for m in _rows(1, hi)), + mxfp8_low_m_gemm, + lambda spec: (((spec["n"], spec["k"]), 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 (), +) + +# 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"]),), +) +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, + ATTN_MEGA, + ATTN_C1, + FUSED_DECODE_ATTN, + DENSE_DECODE_SPLIT, + SPARSE_DECODE_PRODUCER, + PRODUCER, + INDEX_SCORE, + LONG_SELECT, + MC_ALLREDUCE, + ATTN_PUSH, + MOE_PUSH, + MC_FINALIZE, + MOE_ENGINE, + MOE16_ENGINE, + ROUTER_GEMV, + ROUTER_GEMV_FP4, + ROUTER_TC, + MX_GEMV, + MX_GEMV_BF16, + LOW_M_GEMM, + SHARD_ACCEPT, + TARGET_HEAD, + DRAFT_STEP, + 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 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: 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 + + +def router_path(num_tokens: int) -> Optional[str]: + """The router kernel when the MoE runs unfused: the fp32 GEMV, the tensor-core router, or None (built-in gate).""" + if ROUTER_GEMV.covers(num_tokens): + return ROUTER_GEMV.path + if ROUTER_TC.covers(num_tokens): + return ROUTER_TC.path + return None + + +MX_GEMV_PATH = MX_GEMV.path +MX_GEMV_BF16_PATH = MX_GEMV_BF16.path +LOW_M_GEMM_PATH = LOW_M_GEMM.path +# 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``: the projection is marked for the + fused low-M GEMM. + """ + 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 low_m_gemm and num_tokens <= LOW_M_GEMM_MAX_TOKENS.get((n, k), 0): + return LOW_M_GEMM_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 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() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/__init__.py new file mode 100644 index 000000000000..a111efce7a4a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kernels of the MiniMax-M3 NVFP4 / sm_103 / tp4 target.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/long_select.py new file mode 100644 index 000000000000..4fb5064fa981 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/long_select.py @@ -0,0 +1,62 @@ +# 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_select_long``, 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 + + +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, +) -> 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 (``dispatch.SELECT_EARLY_MAX_ROWS``). + """ + 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( + max_score_kv, + n_valid_blocks.contiguous(), + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_lc.py new file mode 100644 index 000000000000..09b4fc4ee57d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_lc.py @@ -0,0 +1,112 @@ +# 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 (``dispatch.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 + + +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, +): + """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``). + """ + 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), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega.py new file mode 100644 index 000000000000..da6a0ddffcca --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 attention megakernel: ``m3_attn_mega``, one launch for a one-request, one-token decode layer. + +The QKV GEMV (bit-identical to ``m3_mx_gemv``: split 4 with contiguous k-tiles, one accumulator, the MXFP8 +activation as produced by the all-reduce, the split-K partials reduced in the owner CTA over DSMEM in rank order +and rounded to bf16), the fused decode attention and the o_proj GEMV run in one persistent grid, joined by per-tile +ready counters instead of a PDL boundary. + +The kernel is ``torch.ops.trtllm.minimax_m3_attn_mega`` (``cpp/tensorrt_llm/kernels/minimaxM3/attnMega.cu``), +compiled for geometry (1, 1, 1) with the MXFP8 hand-off only. + +Grid: 7 clusters of 16 CTAs. The first attention cluster(s) run the attention; the last 80 CTAs are the QKV units +(row tile u / 4, split u % 4; the four splits of a tile are consecutive CTAs of one cluster). Each unit's owner CTA +of a 32-row slice stores it to ``packed_out`` and release-increments the tile's counter; every attention CTA waits +for 4 more increments on all 20 tiles than at its last call (a per-CTA target it alone keeps), then reads the rows. + +One kernel replaces the per-layer chain that follows the packed QKV GEMM in short-sequence decode mode (every +request attends at most 16 KV blocks of 128 tokens, so block selection is the identity): + +* q/k/index-k RMSNorm + NeoX partial RoPE, NVFP4 K/V cache insert and FP8 index-K insert + (``minimaxM3Nvfp4QKVIndexerNormRopeKVInsertKernel``), +* block-sparse GQA decode over the 16 blocks (``_gqa_sparse_decode_kernel``), +* the split-K LSE merge, optionally with the MXFP8 copy of the output (``_merge_topk_attn_out_kernel``). + +Geometry: 16 query heads, one KV head, one index head, head_dim 128, rotary_dim 64, page 128. + +Launch: clusters of ``16 / BPC`` CTAs; CTA ``r`` owns KV blocks ``[r * BPC, r * BPC + BPC)`` for ``QG`` query +tokens of one request (``decode_query_len / QG`` clusters per request). Each block keeps its own softmax partial, +exactly the Triton producer's chunk, so the merge sees the same 16 partials. + +* Before ``griddepcontrol.wait``: request metadata, RoPE rows, norm weights, the packed K chunks (registers) and the + packed V tile (``cp.async``). These are staged before the CUDA graph or written by earlier steps; nothing the QKV + GEMM writes is read, and nothing is written, before the wait. +* After the wait: K is dequantized to fp16 (exact) while the packed-QKV loads are in flight; every CTA + normalizes/ropes the query rows it needs (recomputed rather than exchanged); the CTA whose block holds a new + token's position also normalizes, ropes and quantizes that token's K, V and index-K exactly as the C++ producer + does, and writes the dequantized row into its fp16 tiles. The PDL dependents are released once every K/V byte + has been read: an earlier release lets their weight streams delay the KV reads. +* S^T = K Q^T and O^T = V^T P^T run on ``tcgen05.mma`` (FP16 operands from shared memory, FP32 accumulators in + TMEM). The softmax runs in the TMEM-reading warps with the producer's per-block math (block max, exp2, fp16 P, + log-domain LSE), so each partial equals the Triton producer's chunk partial up to accumulation order. The + producer warps dequantize V meanwhile. +* Partials go from registers to the merging CTA with ``st.async`` + mbarrier transaction counts (DSMEM); CTA ``r`` + merges its heads with the merge kernel's LSE weighting and writes the BF16 output (and the MXFP8 hand-off). + +CUDA C++ rather than Triton or CuTe DSL: the design needs cluster DSMEM pushes, loads issued before the PDL wait, +tcgen05/TMEM, and bit-exact reuse of the C++ producer's norm/RoPE/NVFP4 arithmetic. +""" + +from typing import Optional + +import torch + +NUM_HEADS = 16 +HEAD_DIM = 128 +PAGE_SIZE = 128 +MAX_BLOCKS = 16 +PACKED_WIDTH = (NUM_HEADS + 4) * HEAD_DIM + +_module = None + + +def _load_module(): + """The compiled ops' namespace (``torch.ops.trtllm``), resolved once; ``_module`` stays None until then.""" + global _module + if _module is None: + if not hasattr(torch.ops.trtllm, "minimax_m3_attn_mega"): + raise RuntimeError( + "torch.ops.trtllm.minimax_m3_attn_mega is not registered: this TensorRT-LLM build lacks " + "cpp/tensorrt_llm/kernels/minimaxM3/attnMega.cu" + ) + _module = torch.ops.trtllm + return _module + + +def _check(cond: bool, msg: str) -> None: + if not cond: + raise ValueError(f"m3_fused_decode_attn: {msg}") + + +class MegaWorkspace: + """Per-layer persistent state of the megakernel (allocate before CUDA-graph capture; nothing is ever reset): the + two call-parity buffers of the packed QKV rows and of the attention's MXFP8 hand-off (marker-filled), and + ``flags[32 + CTA]``, each CTA's call counter (its low bit selects the parity).""" + + def __init__(self, device, max_tokens: int = 8): + # Two call-parity buffers of the packed QKV rows, all set to the 0xFFFF "not yet" marker. + self.qkv_out = torch.full( + (2, max_tokens, PACKED_WIDTH), -1, dtype=torch.int16, device=device + ).view(torch.bfloat16) + self.flags = torch.zeros(256, dtype=torch.int32, device=device) + # Two call-parity buffers of the attention's MXFP8 hand-off to the o_proj units, all 0xFF ("not yet"). + self.q8 = torch.full((2 * max_tokens * 2048,), 255, dtype=torch.uint8, device=device) + self.sf = torch.full((2 * 16 * 512,), 255, dtype=torch.uint8, device=device) + self.calls = 0 # host mirror of the call count (eager use only; graph replays advance the device copy) + + +def m3_attn_mega( + xq: torch.Tensor, + xsf: torch.Tensor, + qkv_weight: torch.Tensor, + qkv_weight_scale: torch.Tensor, + ws: MegaWorkspace, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + o_weight: torch.Tensor, + o_weight_scale: torch.Tensor, + *, + eps: float, + sm_scale: float, + decode_query_len: int, + output: torch.Tensor, + y: torch.Tensor, + use_pdl: bool = True, + ar: Optional[tuple] = None, +) -> torch.Tensor: + """QKV projection (bit-identical to m3_mx_gemv) of the MXFP8 layer input ``xq`` / ``xsf``, the fused decode + attention (bit-identical to m3_fused_decode_attn) and the o_proj projection (bit-identical to m3_mx_gemv) into + ``y`` [T, 6144], in one launch. + ``output`` receives the attention rows (bf16). ``ar`` = m3_mx_gemv's (multicast workspace address, flags address, + rank, world, bucket tokens, mode 2): also push y into the next m3_mc_allreduce's exchange. Returns the packed QKV + rows of this call.""" + num_tokens = int(xq.shape[0]) + num_reqs = int(seq_lens.shape[0]) + _check( + num_tokens == num_reqs * decode_query_len and num_tokens <= ws.qkv_out.shape[0], + "token count", + ) + _check( + num_reqs == 1 and decode_query_len == 1, + "the megakernel runs one request with one query token (one " + "attention cluster); use the m3_mx_gemv + m3_fused_decode_attn + m3_mx_gemv path otherwise", + ) + _check( + xq.dtype == torch.float8_e4m3fn and xq.shape[1] == 6144 and xq.stride(1) == 1, + "xq must be e4m3 [T, 6144]", + ) + _check(qkv_weight.shape == (PACKED_WIDTH, 6144), "QKV weight must be [2560, 6144]") + out2d = output.view(num_tokens, NUM_HEADS * HEAD_DIM) + q8, sf = ws.q8, ws.sf + _check( + y.dtype == torch.bfloat16 and y.shape == (num_tokens, 6144) and y.is_contiguous(), + "y must be bf16 [T, 6144]", + ) + qg, bpc = 1, 1 # one 16-CTA attention cluster, one KV block per CTA + _load_module().minimax_m3_attn_mega( + xq.view(torch.uint8), + xsf.view(torch.uint8), + qkv_weight.view(torch.uint8), + qkv_weight_scale.view(torch.uint8), + ws.qkv_out, + ws.flags, + kv_data, + kv_scale, + index_k_cache.view(torch.uint8), + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + float(eps), + float(sm_scale), + int(decode_query_len), + qg, + bpc, + out2d, + q8, + sf, + bool(use_pdl), + o_weight.view(torch.uint8), + o_weight_scale.view(torch.uint8), + y, + *((ar[0], ar[1], ar[2], ar[3], ar[4], 2) if ar is not None else (0, 0, 0, 1, 8, 0)), + ) + out_rows = ws.qkv_out[ws.calls & 1, :num_tokens] + ws.calls += 1 + return out_rows + + +__all__ = ["MegaWorkspace", "m3_attn_mega"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega_c1.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega_c1.py new file mode 100644 index 000000000000..e4898442f227 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega_c1.py @@ -0,0 +1,358 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 attention + o_proj kernel: ``m3_attn_c1``, one launch for a one-request, 6-token verify layer. + +The fused decode attention on the packed QKV rows and the o_proj GEMV run in one 96-CTA launch (6 clusters of 16): +every attention CTA is also an o_proj unit, fed by the attention's MXFP8 hand-off through the workspace (Lamport, +call parity in its flags), and the o_proj rows are pushed into the next multicast all-reduce's exchange. + +The kernel is ``torch.ops.trtllm.minimax_m3_attn_c1`` (``cpp/tensorrt_llm/kernels/minimaxM3/attnC1.cu``), compiled +for ``C1_GEOMETRIES`` only. + +One kernel replaces the per-layer chain that follows the packed QKV GEMM in short-sequence decode mode (every +request attends at most 16 KV blocks of 128 tokens, so block selection is the identity): + +* q/k/index-k RMSNorm + NeoX partial RoPE, NVFP4 K/V cache insert and FP8 index-K insert + (``minimaxM3Nvfp4QKVIndexerNormRopeKVInsertKernel``), +* block-sparse GQA decode over the 16 blocks (``_gqa_sparse_decode_kernel``), +* the split-K LSE merge, optionally with the MXFP8 copy of the output (``_merge_topk_attn_out_kernel``). + +Geometry: 16 query heads, one KV head, one index head, head_dim 128, rotary_dim 64, page 128. + +Launch: clusters of ``16 / BPC`` CTAs; CTA ``r`` owns KV blocks ``[r * BPC, r * BPC + BPC)`` for ``QG`` query +tokens of one request (``decode_query_len / QG`` clusters per request). Each block keeps its own softmax partial, +exactly the Triton producer's chunk, so the merge sees the same 16 partials. + +* Before ``griddepcontrol.wait``: request metadata, RoPE rows, norm weights, the packed K chunks (registers) and the + packed V tile (``cp.async``). These are staged before the CUDA graph or written by earlier steps; nothing the QKV + GEMM writes is read, and nothing is written, before the wait. +* After the wait: K is dequantized to fp16 (exact) while the packed-QKV loads are in flight; every CTA + normalizes/ropes the query rows it needs (recomputed rather than exchanged); the CTA whose block holds a new + token's position also normalizes, ropes and quantizes that token's K, V and index-K exactly as the C++ producer + does, and writes the dequantized row into its fp16 tiles. The PDL dependents are released once every K/V byte + has been read: an earlier release lets their weight streams delay the KV reads. +* S^T = K Q^T and O^T = V^T P^T run on ``tcgen05.mma`` (FP16 operands from shared memory, FP32 accumulators in + TMEM). The softmax runs in the TMEM-reading warps with the producer's per-block math (block max, exp2, fp16 P, + log-domain LSE), so each partial equals the Triton producer's chunk partial up to accumulation order. The + producer warps dequantize V meanwhile. +* Partials go from registers to the merging CTA with ``st.async`` + mbarrier transaction counts (DSMEM); CTA ``r`` + merges its heads with the merge kernel's LSE weighting and writes the BF16 output (and the MXFP8 hand-off). + +CUDA C++ rather than Triton or CuTe DSL: the design needs cluster DSMEM pushes, loads issued before the PDL wait, +tcgen05/TMEM, and bit-exact reuse of the C++ producer's norm/RoPE/NVFP4 arithmetic. +""" + +from typing import Optional + +import torch + +NUM_HEADS = 16 +HEAD_DIM = 128 +PAGE_SIZE = 128 +MAX_BLOCKS = 16 +PACKED_WIDTH = (NUM_HEADS + 4) * HEAD_DIM + +# The geometries of m3_fused_attn, which default_geometry() reads (only C1_GEOMETRIES is compiled into this op). +GEOMETRIES = frozenset( + [ + (1, 1, 1), + (1, 1, 2), + (2, 1, 1), + (2, 1, 2), + (3, 1, 1), + (4, 1, 1), + (5, 1, 1), + (6, 1, 1), + (6, 1, 2), + (7, 1, 1), + (8, 1, 1), + ] +) + +_module = None + + +def _load_module(): + """The compiled ops' namespace (``torch.ops.trtllm``), resolved once; ``_module`` stays None until then.""" + global _module + if _module is None: + if not hasattr(torch.ops.trtllm, "minimax_m3_attn_c1"): + raise RuntimeError( + "torch.ops.trtllm.minimax_m3_attn_c1 is not registered: this TensorRT-LLM build lacks " + "cpp/tensorrt_llm/kernels/minimaxM3/attnC1.cu" + ) + _module = torch.ops.trtllm + return _module + + +def default_geometry(decode_query_len: int, num_requests: int) -> tuple: + """(tokens per cluster, KV blocks per CTA) for a step shape, as m3_fused_attn.default_geometry.""" + clusters = decode_query_len * num_requests + if clusters <= 6: + return (1, 1) + if (decode_query_len, 1, 2) in GEOMETRIES and clusters <= 16: + return (1, 2) + return (1, 1) + + +def _check(cond: bool, msg: str) -> None: + if not cond: + raise ValueError(f"m3_fused_decode_attn: {msg}") + + +def m3_fused_decode_attn_checks( + num_tokens, + num_reqs, + decode_query_len, + kv_data, + kv_scale, + index_k_cache, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + output, +) -> None: + """Argument checks shared by the attention entry points (the packed or fused-projection QKV source).""" + _check( + 1 <= decode_query_len <= 8, f"decode_query_len must be in [1, 8], got {decode_query_len}" + ) + _check( + num_tokens == num_reqs * decode_query_len, + "token rows must equal num_requests * decode_query_len", + ) + _check( + kv_data.dtype == torch.uint8 + and kv_data.dim() == 5 + and tuple(kv_data.shape[1:]) == (2, 1, 128, 64) + and kv_data.stride(4) == 1 + and kv_data.stride(3) == 64 + and kv_data.stride(0) % 16 == 0 + and kv_data.stride(1) % 16 == 0 + and kv_data.data_ptr() % 16 == 0, + "kv_data must be uint8 [pages, 2, 1, 128, 64] with contiguous 16-byte aligned head blocks", + ) + _check( + kv_scale.dtype == torch.uint8 + and kv_scale.dim() == 5 + and tuple(kv_scale.shape[1:]) == (2, 1, 128, 8) + and kv_scale.stride(4) == 1 + and kv_scale.stride(3) == 8 + and kv_scale.stride(0) % 16 == 0 + and kv_scale.stride(1) % 16 == 0 + and kv_scale.data_ptr() % 16 == 0 + and kv_scale.shape[0] == kv_data.shape[0], + "kv_scale must be uint8 [pages, 2, 1, 128, 8] with contiguous 16-byte aligned head blocks", + ) + _check( + 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(0) % 4 == 0 + and index_k_cache.shape[0] == kv_data.shape[0], + "index_k_cache must be fp8 [pages, 1, 128, 128]", + ) + for name, t in ( + ("out_cache_loc", out_cache_loc), + ("position_ids", position_ids), + ("block_table", block_table), + ("seq_lens", seq_lens), + ): + _check(t.dtype == torch.int32 and t.is_cuda, f"{name} must be CUDA int32") + _check( + out_cache_loc.is_contiguous() and out_cache_loc.numel() >= num_tokens, + "out_cache_loc too short", + ) + _check( + position_ids.is_contiguous() and position_ids.numel() == num_tokens, + "position_ids must have T entries", + ) + _check( + seq_lens.is_contiguous() + and block_table.dim() == 2 + and block_table.stride(1) == 1 + and block_table.shape[0] >= num_reqs, + "block_table must be [num_requests, max_blocks] row-major", + ) + for name, t in ( + ("kv_quant_scale", kv_quant_scale), + ("k_global_scale", k_global_scale), + ("v_global_scale", v_global_scale), + ): + _check( + t.dtype == torch.float32 and t.is_cuda and t.is_contiguous(), + f"{name} must be CUDA fp32", + ) + _check(kv_quant_scale.numel() >= 3, "kv_quant_scale must hold [Q, K, V]") + for name, t in ( + ("q_norm_weight", q_norm_weight), + ("k_norm_weight", k_norm_weight), + ("index_k_norm_weight", index_k_norm_weight), + ): + _check( + t.dtype == torch.bfloat16 + and t.numel() == HEAD_DIM + and t.is_contiguous() + and t.data_ptr() % 8 == 0, + f"{name} must be contiguous bf16 [128]", + ) + _check( + rotary_cos_sin.dtype == torch.float32 + and rotary_cos_sin.is_contiguous() + and rotary_cos_sin.dim() == 3 + and tuple(rotary_cos_sin.shape[1:]) == (2, 32), + "rotary_cos_sin must be fp32 [max_pos, 2, 32]", + ) + out2d = output.view(num_tokens, NUM_HEADS * HEAD_DIM) + _check( + output.dtype == torch.bfloat16 and out2d.stride(1) == 1 and out2d.data_ptr() % 8 == 0, + "output must be bf16 [T, 2048] with contiguous rows", + ) + + +class MegaWorkspace: + """Per-layer persistent state of the megakernel (allocate before CUDA-graph capture; nothing is ever reset): the + two call-parity buffers of the packed QKV rows and of the attention's MXFP8 hand-off (marker-filled), and + ``flags[32 + CTA]``, each CTA's call counter (its low bit selects the parity).""" + + def __init__(self, device, max_tokens: int = 8): + # Two call-parity buffers of the packed QKV rows, all set to the 0xFFFF "not yet" marker. + self.qkv_out = torch.full( + (2, max_tokens, PACKED_WIDTH), -1, dtype=torch.int16, device=device + ).view(torch.bfloat16) + self.flags = torch.zeros(256, dtype=torch.int32, device=device) + # Two call-parity buffers of the attention's MXFP8 hand-off to the o_proj units, all 0xFF ("not yet"). + self.q8 = torch.full((2 * max_tokens * 2048,), 255, dtype=torch.uint8, device=device) + self.sf = torch.full((2 * 16 * 512,), 255, dtype=torch.uint8, device=device) + self.calls = 0 # host mirror of the call count (eager use only; graph replays advance the device copy) + + +C1_GEOMETRIES = frozenset([(6, 1, 1)]) + + +def c1_applies(decode_query_len: int, num_requests: int) -> bool: + """Whether the step shape has a c1 kernel: 96 attention CTAs (6 clusters of 16), one o_proj unit each.""" + qg, bpc = default_geometry(decode_query_len, num_requests) + return ( + num_requests * (decode_query_len // qg) * (16 // bpc) == 96 + and (decode_query_len, qg, bpc) in C1_GEOMETRIES + ) + + +def m3_attn_c1( + packed: torch.Tensor, + ws: MegaWorkspace, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + o_weight: torch.Tensor, + o_weight_scale: torch.Tensor, + *, + eps: float, + sm_scale: float, + decode_query_len: int, + output: torch.Tensor, + y: torch.Tensor, + use_pdl: bool = True, + ar: Optional[tuple] = None, +) -> None: + """The fused decode attention on the packed QKV rows (bit-identical to m3_fused_decode_attn) and the o_proj + projection (bit-identical to m3_mx_gemv) into ``y`` [T, 6144], in one 96-CTA launch: every attention CTA is also + an o_proj unit, fed by the attention's MXFP8 hand-off through ``ws`` (Lamport, call parity in ``ws.flags``). + ``ar`` = m3_mx_gemv's (multicast workspace address, flags address, rank, world, bucket tokens): also push y into + the next m3_mc_allreduce's exchange.""" + num_tokens = int(packed.shape[0]) + num_reqs = int(seq_lens.shape[0]) + _check( + packed.dtype == torch.bfloat16 + and packed.dim() == 2 + and packed.shape[1] == PACKED_WIDTH + and packed.stride(1) == 1 + and packed.stride(0) % 4 == 0 + and packed.data_ptr() % 8 == 0, + "packed must be bf16 [T, 2560] with 8-byte aligned contiguous rows", + ) + m3_fused_decode_attn_checks( + num_tokens, + num_reqs, + decode_query_len, + kv_data, + kv_scale, + index_k_cache, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + output, + ) + qg, bpc = default_geometry(decode_query_len, num_reqs) + _check( + c1_applies(decode_query_len, num_reqs), f"no c1 kernel for {num_reqs} x {decode_query_len}" + ) + _check( + y.dtype == torch.bfloat16 and y.shape == (num_tokens, 6144) and y.is_contiguous(), + "y must be bf16 [T, 6144]", + ) + out2d = output.view(num_tokens, NUM_HEADS * HEAD_DIM) + _load_module().minimax_m3_attn_c1( + packed, + ws.flags, + kv_data, + kv_scale, + index_k_cache.view(torch.uint8), + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + float(eps), + float(sm_scale), + int(decode_query_len), + qg, + bpc, + out2d, + ws.q8, + ws.sf, + bool(use_pdl), + o_weight.view(torch.uint8), + o_weight_scale.view(torch.uint8), + y, + *((ar[0], ar[1], ar[2], ar[3], ar[4], 2) if ar is not None else (0, 0, 0, 1, 8, 0)), + ) + ws.calls += 1 + + +__all__ = ["MegaWorkspace", "c1_applies", "m3_attn_c1"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_draft_step.py new file mode 100644 index 000000000000..f15d8dbeff64 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_draft_step.py @@ -0,0 +1,179 @@ +# 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 .. import dispatch +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 .m3_shard_accept import AcceptWorkspace, create_workspace + + +@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 + _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 + ``dispatch.DRAFT_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 dispatch.DRAFT_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) + return DraftHead(weight=weight, wf=wf, vocab_offset=offset) + + +_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(): + dispatch.capture_missed(dispatch.DRAFT_STEP.path, 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 + _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, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_fused_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_fused_attn.py new file mode 100644 index 000000000000..c404166fb320 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_fused_attn.py @@ -0,0 +1,349 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 fused short-sequence decode attention: ``m3_fused_decode_attn``. + +The kernel is ``torch.ops.trtllm.minimax_m3_fused_decode_attn`` (``cpp/tensorrt_llm/kernels/minimaxM3/ +fusedDecodeAttn.cu``); it is compiled for the geometries in ``GEOMETRIES``, each with the MXFP8 hand-off. + +One kernel replaces the per-layer chain that follows the packed QKV GEMM in short-sequence decode mode (every +request attends at most 16 KV blocks of 128 tokens, so block selection is the identity): + +* q/k/index-k RMSNorm + NeoX partial RoPE, NVFP4 K/V cache insert and FP8 index-K insert + (``minimaxM3Nvfp4QKVIndexerNormRopeKVInsertKernel``), +* block-sparse GQA decode over the 16 blocks (``_gqa_sparse_decode_kernel``), +* the split-K LSE merge, optionally with the MXFP8 copy of the output (``_merge_topk_attn_out_kernel``). + +Geometry: 16 query heads, one KV head, one index head, head_dim 128, rotary_dim 64, page 128. + +Launch: clusters of ``16 / BPC`` CTAs; CTA ``r`` owns KV blocks ``[r * BPC, r * BPC + BPC)`` for ``QG`` query +tokens of one request (``decode_query_len / QG`` clusters per request). Each block keeps its own softmax partial, +exactly the Triton producer's chunk, so the merge sees the same 16 partials. + +* Before ``griddepcontrol.wait``: request metadata, RoPE rows, norm weights, the packed K chunks (registers) and the + packed V tile (``cp.async``). These are staged before the CUDA graph or written by earlier steps; nothing the QKV + GEMM writes is read, and nothing is written, before the wait. +* After the wait: K is dequantized to fp16 (exact) while the packed-QKV loads are in flight; every CTA + normalizes/ropes the query rows it needs (recomputed rather than exchanged); the CTA whose block holds a new + token's position also normalizes, ropes and quantizes that token's K, V and index-K exactly as the C++ producer + does, and writes the dequantized row into its fp16 tiles. The PDL dependents are released once every K/V byte + has been read: an earlier release lets their weight streams delay the KV reads. +* S^T = K Q^T and O^T = V^T P^T run on ``tcgen05.mma`` (FP16 operands from shared memory, FP32 accumulators in + TMEM). The softmax runs in the TMEM-reading warps with the producer's per-block math (block max, exp2, fp16 P, + log-domain LSE), so each partial equals the Triton producer's chunk partial up to accumulation order. The + producer warps dequantize V meanwhile. +* Partials go from registers to the merging CTA with ``st.async`` + mbarrier transaction counts (DSMEM); CTA ``r`` + merges its heads with the merge kernel's LSE weighting and writes the BF16 output (and the MXFP8 hand-off). + +CUDA C++ rather than Triton or CuTe DSL: the design needs cluster DSMEM pushes, loads issued before the PDL wait, +tcgen05/TMEM, and bit-exact reuse of the C++ producer's norm/RoPE/NVFP4 arithmetic. +""" + +from typing import Optional + +import torch + +NUM_HEADS = 16 +HEAD_DIM = 128 +PAGE_SIZE = 128 +MAX_BLOCKS = 16 +PACKED_WIDTH = (NUM_HEADS + 4) * HEAD_DIM + +# (decode_query_len, tokens per cluster, KV blocks per CTA) geometries compiled into the op: what +# default_geometry() returns for every step shape the target routes here (at most 16 tokens). +GEOMETRIES = frozenset( + [ + (1, 1, 1), + (1, 1, 2), + (2, 1, 1), + (2, 1, 2), + (3, 1, 1), + (4, 1, 1), + (5, 1, 1), + (6, 1, 1), + (6, 1, 2), + (7, 1, 1), + (8, 1, 1), + ] +) + +_module = None + + +def _load_module(): + """The compiled ops' namespace (``torch.ops.trtllm``), resolved once; ``_module`` stays None until then.""" + global _module + if _module is None: + if not hasattr(torch.ops.trtllm, "minimax_m3_fused_decode_attn"): + raise RuntimeError( + "torch.ops.trtllm.minimax_m3_fused_decode_attn is not registered: this TensorRT-LLM " + "build lacks cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.cu" + ) + _module = torch.ops.trtllm + return _module + + +def default_geometry(decode_query_len: int, num_requests: int) -> tuple: + """(tokens per cluster, KV blocks per CTA) for a step shape: one token per 16-CTA cluster up to 6 clusters, + then KV blocks paired into 8-CTA clusters up to 16 clusters where that geometry is compiled.""" + clusters = decode_query_len * num_requests + if clusters <= 6: + return (1, 1) + if (decode_query_len, 1, 2) in GEOMETRIES and clusters <= 16: + return (1, 2) + return (1, 1) + + +def _check(cond: bool, msg: str) -> None: + if not cond: + raise ValueError(f"m3_fused_decode_attn: {msg}") + + +def m3_fused_decode_attn_checks( + num_tokens, + num_reqs, + decode_query_len, + kv_data, + kv_scale, + index_k_cache, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + output, +) -> None: + """Argument checks shared by the attention entry points (the packed or fused-projection QKV source).""" + _check( + 1 <= decode_query_len <= 8, f"decode_query_len must be in [1, 8], got {decode_query_len}" + ) + _check( + num_tokens == num_reqs * decode_query_len, + "token rows must equal num_requests * decode_query_len", + ) + _check( + kv_data.dtype == torch.uint8 + and kv_data.dim() == 5 + and tuple(kv_data.shape[1:]) == (2, 1, 128, 64) + and kv_data.stride(4) == 1 + and kv_data.stride(3) == 64 + and kv_data.stride(0) % 16 == 0 + and kv_data.stride(1) % 16 == 0 + and kv_data.data_ptr() % 16 == 0, + "kv_data must be uint8 [pages, 2, 1, 128, 64] with contiguous 16-byte aligned head blocks", + ) + _check( + kv_scale.dtype == torch.uint8 + and kv_scale.dim() == 5 + and tuple(kv_scale.shape[1:]) == (2, 1, 128, 8) + and kv_scale.stride(4) == 1 + and kv_scale.stride(3) == 8 + and kv_scale.stride(0) % 16 == 0 + and kv_scale.stride(1) % 16 == 0 + and kv_scale.data_ptr() % 16 == 0 + and kv_scale.shape[0] == kv_data.shape[0], + "kv_scale must be uint8 [pages, 2, 1, 128, 8] with contiguous 16-byte aligned head blocks", + ) + _check( + 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(0) % 4 == 0 + and index_k_cache.shape[0] == kv_data.shape[0], + "index_k_cache must be fp8 [pages, 1, 128, 128]", + ) + for name, t in ( + ("out_cache_loc", out_cache_loc), + ("position_ids", position_ids), + ("block_table", block_table), + ("seq_lens", seq_lens), + ): + _check(t.dtype == torch.int32 and t.is_cuda, f"{name} must be CUDA int32") + _check( + out_cache_loc.is_contiguous() and out_cache_loc.numel() >= num_tokens, + "out_cache_loc too short", + ) + _check( + position_ids.is_contiguous() and position_ids.numel() == num_tokens, + "position_ids must have T entries", + ) + _check( + seq_lens.is_contiguous() + and block_table.dim() == 2 + and block_table.stride(1) == 1 + and block_table.shape[0] >= num_reqs, + "block_table must be [num_requests, max_blocks] row-major", + ) + for name, t in ( + ("kv_quant_scale", kv_quant_scale), + ("k_global_scale", k_global_scale), + ("v_global_scale", v_global_scale), + ): + _check( + t.dtype == torch.float32 and t.is_cuda and t.is_contiguous(), + f"{name} must be CUDA fp32", + ) + _check(kv_quant_scale.numel() >= 3, "kv_quant_scale must hold [Q, K, V]") + for name, t in ( + ("q_norm_weight", q_norm_weight), + ("k_norm_weight", k_norm_weight), + ("index_k_norm_weight", index_k_norm_weight), + ): + _check( + t.dtype == torch.bfloat16 + and t.numel() == HEAD_DIM + and t.is_contiguous() + and t.data_ptr() % 8 == 0, + f"{name} must be contiguous bf16 [128]", + ) + _check( + rotary_cos_sin.dtype == torch.float32 + and rotary_cos_sin.is_contiguous() + and rotary_cos_sin.dim() == 3 + and tuple(rotary_cos_sin.shape[1:]) == (2, 32), + "rotary_cos_sin must be fp32 [max_pos, 2, 32]", + ) + out2d = output.view(num_tokens, NUM_HEADS * HEAD_DIM) + _check( + output.dtype == torch.bfloat16 and out2d.stride(1) == 1 and out2d.data_ptr() % 8 == 0, + "output must be bf16 [T, 2048] with contiguous rows", + ) + + +def m3_fused_decode_attn( + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + *, + eps: float, + sm_scale: float, + decode_query_len: int, + output: torch.Tensor, + mxfp8_out: Optional[tuple] = None, + use_pdl: bool = True, + geometry: Optional[tuple] = None, +) -> None: + """Fused norm/RoPE/KV-insert + short-sequence sparse decode + split-K merge, written into ``output``. + + packed: [T, 2560] bf16 QKV+index GEMM output, T = num_requests * decode_query_len. + kv_data: uint8 [num_pages, 2, 1, 128, 64] HND NVFP4 pool (K plane 0, V plane 1); written. + kv_scale: uint8 [num_pages, 2, 1, 128, 8] E4M3 block scales (K linear, V 4x4 token quads); written. + index_k_cache: float8_e4m3fn [num_pages, 1, 128, 128]; written. + out_cache_loc, position_ids: int32 [>= T] / [T]; block_table int32 [num_requests, >= 1]; seq_lens int32 + [num_requests] (attended length including this step's tokens, at most 16 * 128). + kv_quant_scale: fp32 [3] (Q, K, V) quantization scales; k/v_global_scale: fp32 [1] dequantization scales. + output: bf16 [T, 2048] (or [T, 16, 128], contiguous rows). + mxfp8_out: (e4m3 [>= T, 2048], uint8 swizzled UE8M0 scales) hand-off for the o_proj GEMM (required: only + the MXFP8 variants are compiled), + same layout as the merge kernel's MERGE_MXFP8 output. + geometry: (query tokens per cluster, KV blocks per CTA); default_geometry() when None. + """ + num_tokens = int(packed.shape[0]) + num_reqs = int(seq_lens.shape[0]) + _check( + 1 <= decode_query_len <= 8, f"decode_query_len must be in [1, 8], got {decode_query_len}" + ) + _check( + num_tokens == num_reqs * decode_query_len, + "packed rows must equal num_requests * decode_query_len", + ) + _check( + packed.dtype == torch.bfloat16 + and packed.dim() == 2 + and packed.shape[1] == PACKED_WIDTH + and packed.stride(1) == 1 + and packed.stride(0) % 4 == 0 + and packed.data_ptr() % 8 == 0, + "packed must be bf16 [T, 2560] with 8-byte aligned contiguous rows", + ) + m3_fused_decode_attn_checks( + num_tokens, + num_reqs, + decode_query_len, + kv_data, + kv_scale, + index_k_cache, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + output, + ) + out2d = output.view(num_tokens, NUM_HEADS * HEAD_DIM) + _check( + mxfp8_out is not None, + "only the variants with the MXFP8 hand-off are compiled; pass mxfp8_out", + ) + q8, sf = mxfp8_out + _check( + q8.dtype == torch.float8_e4m3fn + and q8.dim() == 2 + and q8.shape[1] == NUM_HEADS * HEAD_DIM + and q8.shape[0] >= num_tokens + and q8.is_contiguous(), + "mxfp8 data must be e4m3 [>= T, 2048]", + ) + _check(sf.dtype == torch.uint8 and sf.is_contiguous(), "mxfp8 scales must be contiguous uint8") + q8 = q8.view(torch.uint8) + if num_tokens == 0: + return + qg, bpc = geometry if geometry is not None else default_geometry(decode_query_len, num_reqs) + _check( + (decode_query_len, qg, bpc) in GEOMETRIES, + f"geometry {(decode_query_len, qg, bpc)} is not compiled", + ) + _load_module().minimax_m3_fused_decode_attn( + packed, + kv_data, + kv_scale, + index_k_cache.view(torch.uint8), + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + float(eps), + float(sm_scale), + int(decode_query_len), + qg, + bpc, + out2d, + q8, + sf, + bool(use_pdl), + ) + + +__all__ = ["GEOMETRIES", "default_geometry", "m3_fused_decode_attn"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_index_score.py new file mode 100644 index 000000000000..9748bd1fe4a5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_index_score.py @@ -0,0 +1,645 @@ +# 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/cute_dsl_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.cute_dsl_kernels.blackwell.cute_ptx_utils import ( + EVICT_FIRST, + fp8x4_to_fp16x4, + mma_sync, + simple_tma_copy, +) +from tensorrt_llm._torch.cute_dsl_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 + SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K + # BLOCK_Q must fit one warp's worth of epilogue lanes. + MAX_BLOCK_Q = 32 + 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.cute_dsl_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.cute_dsl_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_103__tp4/kernels/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_mc_allreduce.py new file mode 100644 index 000000000000..bdb207b7a8ef --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_mc_allreduce.py @@ -0,0 +1,368 @@ +# 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 +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, +): + """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. + + 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 + 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), + ) + 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, +): + """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.""" + 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_103__tp4/kernels/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/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_103__tp4/kernels/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_103__tp4/kernels/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_router_gemv.py new file mode 100644 index 000000000000..88a475f509d5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/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 + + +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_103__tp4/kernels/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_sampled_accept.py new file mode 100644 index 000000000000..315e80947213 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_sampled_accept.py @@ -0,0 +1,158 @@ +# 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 dispatch + +# 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 + + +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_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(): + dispatch.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] > 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) + 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_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), + ) + return accepted, num_accepted, gather_ids diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_shard_accept.py new file mode 100644 index 000000000000..e1dafd292d6d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/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 dispatch + +#: Verify rows per call, and exchange rows per call half of the workspace. +SHARD_ROWS = 48 + + +@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(): + dispatch.capture_missed(dispatch.SHARD_ACCEPT.path, 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_103__tp4/kernels/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/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_103__tp4/kernels/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_103__tp4/kernels/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/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_103__tp4/kernels/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_103__tp4/kernels/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/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_103__tp4/kernels/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_103__tp4/kernels/mxfp8_low_m_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mxfp8_low_m_gemm.py new file mode 100644 index 000000000000..04e21b99471e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mxfp8_low_m_gemm.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused MXFP8 activation quantization + low-M MXFP8 GEMM for decode (SM100+). + +``out[M, N] = mxfp8(x)[M, K] @ mxfp8(W)[N, K]^T`` for ``1 <= M <= 8`` in one kernel +(``torch.ops.trtllm.minimax_m3_mxfp8_low_m_gemm``, reached through the ``ops/gemm`` entry). + +* Each CTA owns ``16 * ROW_TILES`` weight rows and ``K / CLUSTER_K`` columns. + Before ``griddepcontrol.wait`` it issues ``cp.async.bulk`` copies of its weight + tile and UE8M0 scale chunks (existing swizzled 128x4 layout) into shared memory, + so the weight stream overlaps the producer kernel under PDL. Configurations + keep the whole grid resident in one wave so every CTA prefetches. +* Activation quantization reproduces ``trtllm::mxfp8_quantize`` with identical + intrinsics (block 32, UE8M0 rounded toward +inf, saturating RN E4M3); each + thread issues all of its loads before converting, one L2 round trip per CTA. +* One ``mma.sync.m16n8k32`` E4M3 per 32-wide scale block and row tile with exact + power-of-two rescaling and FP32 accumulation. Warps reduce through shared + memory (reusing the consumed weight tile); split-K peers push their partial + tiles into the cluster root with ``st.async`` and exit, and the root sums the + ranks in order. +""" + +from typing import Optional + +import torch + +from tensorrt_llm.logger import logger + +from ..ops.gemm.mxfp8_low_m_gemm import mxfp8_low_m_gemm as _low_m_gemm_op + +# (N, K) -> (warps, cluster_k, max_m, row_tiles); measured on GB300 TP4 M3 shapes. The fused kernel beats native +# quantize + GEMM only up to max_m rows; larger M, and every other shape, keeps the native path. These four +# configurations are the ones built into ``torch.ops.trtllm.minimax_m3_mxfp8_low_m_gemm``. +_CONFIGS = { + (2560, 6144): (4, 4, 8, 2), # QKV (q, k, v, index q / k) + (6144, 2048): (8, 2, 8, 4), # o_proj + (6144, 6144): (4, 4, 1, 1), # dense MLP gate_up + (6144, 3072): (4, 2, 8, 2), # dense MLP down +} +# Opt-in shared memory per CTA on SM100 (227 KiB) less a margin for the static mbarriers. +_MAX_DYNAMIC_SMEM = 227 * 1024 - 1024 + +_logged_shapes = set() + + +def _smem_bytes(m: int, k: int, warps: int, cluster_k: int, row_tiles: int) -> int: + """Dynamic shared memory of one CTA, mirroring the host-side size in ``launch``.""" + kc = k // cluster_k + kc_blocks = kc // 32 + tile = 16 * row_tiles * 8 * (6 if m > 32 else 4 if m > 16 else 2 if m > 8 else 1) + w_bytes = 16 * row_tiles * (kc + 32) + red_bytes = warps * tile * 4 + return ( + w_bytes + + row_tiles * (kc_blocks // 4) * 256 + + m * kc + + (m * kc_blocks + 15) // 16 * 16 + + (cluster_k - 1) * tile * 4 + + (0 if red_bytes <= w_bytes else red_bytes) + ) + + +def get_config(n: int, k: int) -> Optional[tuple]: + """Return (warps, cluster_k, max_m, row_tiles) for an (N, K) weight, or None when no configuration is built.""" + return _CONFIGS.get((n, k)) + + +def mxfp8_low_m_gemm( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + use_pdl: bool = False, + config: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """Fused MXFP8 quantize + GEMM; ``weight_scale`` is the swizzled 128x4 buffer. + + Returns None for unsupported operands so callers keep the native path. + """ + if ( + output_dtype != torch.bfloat16 + or x.dtype != torch.bfloat16 + or torch.cuda.get_device_capability(x.device)[0] < 10 + ): + return None + n, k = weight.shape + config = config or get_config(n, k) + if config is None or config not in _CONFIGS.values(): + return None + warps, cluster_k, max_m, row_tiles = config + if ( + x.shape[0] > max_m + or n % (16 * row_tiles) != 0 + or k % (cluster_k * 128) != 0 + or (k // cluster_k // 32) % warps != 0 + or x.data_ptr() % 16 != 0 + or x.stride(0) % 8 != 0 + or weight.data_ptr() % 16 != 0 + or weight_scale.data_ptr() % 16 != 0 + or _smem_bytes(x.shape[0], k, warps, cluster_k, row_tiles) > _MAX_DYNAMIC_SMEM + ): + return None + if (n, k) not in _logged_shapes: + _logged_shapes.add((n, k)) + logger.info( + f"[MXFP8 low-M] fused decode GEMM engaged: N={n} K={k} " + f"warps={warps} cluster_k={cluster_k} row_tiles={row_tiles} pdl={use_pdl}" + ) + return _low_m_gemm_op(x, weight, weight_scale, warps, cluster_k, row_tiles, use_pdl) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/triton_sparse_decode.py new file mode 100644 index 000000000000..c1e0f4332b88 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/triton_sparse_decode.py @@ -0,0 +1,1651 @@ +# 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_103 / 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 + +# One sparse block is exactly one KV page. +SPARSE_BLOCK_SIZE = 128 + +_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 + +# 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 + + +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 + + 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, + 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 +): + 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): + 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. + page = tl.load(bt_row + blk).to(tl.int64) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len + 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, + 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 +): + 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 + 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 _ in tl.range(0, num_candidates): + if QUERY_GROUP_SIZE == 2: + if chunk_equal: + has_block = True + 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 + 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 + 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. + page = tl.load(bt_row + blk).to(tl.int64) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len_max + q_pos_mask = blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None] + 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 both datacenter Blackwell targets. +_SM100F_CAPABILITIES = ((10, 0), (10, 3)) +# 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 +# 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 + 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, + ) + if mxfp8 is not None: + m3_prequant.SLOT.publish( + output.view(total_q, num_heads * head_dim), mxfp8[0][:total_q], mxfp8[1] + ) + + +@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; " + 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 not identity_blocks + 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, + ) + 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) + ) + if kv_nvfp4 and total_q * num_kv_heads <= 64: + # GB300 TP4 Eagle verify (1 KV head, group 16, 6 query tokens): one block per program is fastest up to + # batch 8 (lowm5/bench_decode_chunks.py: b8 11.7 vs 14.8 us with the default split of 4). + num_topk_chunks = 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 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 + ) + 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), + *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, + **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_103__tp4/kernels/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/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_103__tp4/kernels/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_103__tp4/layers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/layers.py new file mode 100644 index 000000000000..3bfafc82cee6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/layers.py @@ -0,0 +1,1682 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 NVFP4 / sm_103 / tp4: the decoder layers, on the built-in modules. + +The attention and MoE modules subclass the built-in ``MiniMaxM3Attention`` and +``MiniMaxM3MoE``: the generic path (prefill, mixed steps, wide steps) runs their +code unchanged, and the checkpoint loads through the built-in loader. What this +target adds is the decode-step kernels, each taken only by the step shapes it +was written for and measured on (GB300, tp4, 1-8 requests): + +* attention: for a short-sequence generation step of an NVFP4 sparse layer (every + request within 16 blocks, so the block selection is the identity) the packed + QKV GEMM is followed by one fused norm/RoPE/KV-insert/sparse-decode/merge + kernel; one request decoding one token runs QKV GEMV + attention + o_proj as + one launch, and one request verifying a draft runs attention + o_proj as one; +* MoE: the router (an fp32 GEMV or a tensor-core router fused with the routed + experts' NVFP4 activation quantization) and the routed experts as one engine + launch for 1-8 tokens, with the shared expert folded in; a second engine + covers 9-13 rows; +* the layer boundaries: a multicast all-reduce fused with the residual add and + the next RMSNorm up to 48 tokens. + +The decoder layer and the text model are this target's own: they construct the +subclassed modules. Their generic forward is the built-in one. +""" + +from __future__ import annotations + +from typing import Callable, Dict, Optional, Tuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +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.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +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_minimaxm3 import ( + MiniMaxM3Attention, + MiniMaxM3MoE, + _build_swiglu_oai_dense_mlp, + _nvfp4_cache_supports_fused_write, + _validate_sparse_attention_runtime_config, + get_moe_layer_ids, + get_sparse_disable_index_value_layer_ids, + get_sparse_layer_ids, + maybe_bcg_minimax_m3_attn_custom_op_inplace, +) +from tensorrt_llm._torch.models.modeling_utils import DecoderModel +from tensorrt_llm._torch.modules.decoder_layer import DecoderLayer +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.modules.linear import ( + flashinfer_mxfp8_autotune_active, + flashinfer_mxfp8_decode_graph_capture_active, +) +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph +from tensorrt_llm._torch.speculative import SpecMetadata +from tensorrt_llm._torch.utils import ( + AuxStreamType, + EventType, + Fp4QuantizedTensor, + is_torch_compiling, +) +from tensorrt_llm.logger import logger + +from . import dispatch +from .dispatch import ( + FUSED_ATTN_MAX_QUERY_LEN, + FUSED_ATTN_MAX_TOKENS, + MOE16_MAX_TOKENS, + MOE16_MIN_TOKENS, + MOE_ENGINE_MAX_TOKENS, + PRODUCER_EARLY_TRIGGER, + ROUTER_TC_MAX_TOKENS, +) +from .kernels import m3_mc_allreduce, m3_prequant +from .kernels.m3_router_gemv import m3_router_gemv, m3_router_gemv_fp4, m3_router_tc +from .linears import PUSH_ATTN, PUSH_MOE, mark_m3_linear +from .ops.attention.m3_index_score import SPLIT_K, m3_index_score +from .ops.attention.m3_producer import m3_producer +from .ops.moe.m3_moe_engine import m3_moe_engine +from .spec_worker import capture_hidden_states + +# The row ranges of the engines, the router and the fused attention are in dispatch.py. +# The shared expert folds into the routed engine from one token (m3_moe with dedicated shared-expert CTAs). +MOE_FOLD_MIN_TOKENS = 1 +# SMs left to the concurrent shared-expert kernels while a routed engine without the shared fold runs (persistent). +MOE_ENGINE_RESERVED_SMS = 24 +# CTAs per request of the CuTe DSL index scorer at every batch (each block is scored once, so the scores do not depend +# on it), and the request counts up to which it compiles the early-issue schedule (kernels/m3_index_score). +INDEX_SCORE_SPLIT_K = SPLIT_K +INDEX_SCORE_EARLY_MAX_REQUESTS = dispatch.INDEX_SCORE_EARLY_MAX_REQUESTS +# The generic layer's boundaries (prefill, Eagle3 capture layers): the low-latency multicast all-reduce + residual + +# RMSNorm beats the native one-shot all-reduce up to 8 tokens; with the shared + routed add folded in, up to its own +# maximum. +MC_AR_MAX_TOKENS = 8 + +# 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] = {} +# The CuTe DSL index scorer of the sparse decode (its kernels are compiled by prebuild_index_scorer). +INDEX_SCORE_PATH = dispatch.INDEX_SCORE.path + + +def _scorer_runner(): + from .kernels.m3_index_score import IndexScoreRunner + + return IndexScoreRunner + + +def _scorer_key( + dtype: torch.dtype, num_heads: int, query_len: int, head_dim: int, split_k: int, early: bool +) -> tuple: + """The scorer's compile key (``IndexScoreRunner``).""" + return (dtype, int(num_heads), int(query_len), int(head_dim), int(split_k), bool(early)) + + +def scorer_variants(layers, query_lens) -> list: + """The scorer keys a decode step of these layers can take: every sparse layer's (index dtype, index heads per + rank, index dim) at each decode query length, with the early-issue schedule (up to + ``INDEX_SCORE_EARLY_MAX_REQUESTS`` requests) and without it, where the kernel serves the geometry.""" + runner = _scorer_runner() + keys = [] + for layer in layers: + attn = getattr(layer, "self_attn", None) + if ( + attn is None + or not getattr(attn, "is_sparse_attention_layer", False) + or not isinstance(attn.attn, MiniMaxM3MsaSparseAttention) + ): + continue + config = attn.attn.m3_config + dtype = torch.float8_e4m3fn if attn.attn.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 = _scorer_key( + dtype, + config.num_index_heads, + qlen, + config.sparse_index_dim, + INDEX_SCORE_SPLIT_K, + early, + ) + if key not in keys: + keys.append(key) + return keys + + +def prebuild_index_scorer(layers, query_lens, device) -> list: + """Compile every key of ``scorer_variants`` 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 = _scorer_runner() + keys = scorer_variants(layers, 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 + + +_EXPERT_TENSOR_NAMES = ( + "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", +) + + +def _nvfp4_swiglu_oai_experts(experts: nn.Module): + """The routed experts' operands the engines read, or raise if the backend is not TRTLLM-Gen NVFP4 SwiGLU-OAI.""" + backend = getattr(experts, "backend", experts) + tensors = [getattr(backend, name, None) for name in _EXPERT_TENSOR_NAMES] + scale = getattr(backend, "fc31_input_scale", None) + if ( + any(t is None for t in tensors) + or scale is None + or getattr(backend, "scaling_vector_size", 16) != 16 + or getattr(backend, "fc31_act_scale", None) is not None + or not type(backend).quantize_input.__qualname__.startswith("TRTLLMGenNvfp4Base") + ): + raise ValueError("experts backend is not TRTLLM-Gen NVFP4 SwiGLU-OAI") + tensors = [t.data if isinstance(t, nn.Parameter) else t for t in tensors] + scale = scale.data if isinstance(scale, nn.Parameter) else scale + return tensors, scale + + +def _dense_mlp(model_config: ModelConfig, intermediate_size: int, layer_idx: int) -> nn.Module: + """The built-in dense SwiGLU-OAI MLP, its projections on this target's decode GEMM dispatch.""" + mlp = _build_swiglu_oai_dense_mlp( + model_config=model_config, intermediate_size=intermediate_size, layer_idx=layer_idx + ) + mark_m3_linear(mlp.gate_up_proj, low_m_gemm=True) + mark_m3_linear(mlp.down_proj, low_m_gemm=True) + return mlp + + +class M3MoE(MiniMaxM3MoE): + """MiniMax-M3 MoE with the decode-step router and expert engines. + + Engines are built per (token count, fusion) by ``prebuild_engines`` from + ``build_targets``, before the executor's warmup (building compiles and warms + the kernel); a forward only looks them up, and a step that finds no engine + is a miss and runs the built-in path. + """ + + def __init__( + self, + model_config: ModelConfig, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + layer_idx: int, + ) -> None: + super().__init__(model_config, aux_stream_dict, layer_idx) + self.layer_idx = layer_idx + if self.shared_experts is not None: + mark_m3_linear(self.shared_experts.gate_up_proj, low_m_gemm=False) + mark_m3_linear(self.shared_experts.down_proj, low_m_gemm=False, push=PUSH_MOE) + # The routed experts' NVFP4 input scale once resolved; False when they cannot take the fused router. + self._fp4_input_scale: Optional[torch.Tensor | bool] = None + self._engines: Dict[Tuple[int, bool, bool], Optional[object]] = {} + self._moe16_engines: Dict[int, Optional[object]] = {} + + def _router_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """The fp32 router logits: the fp32 GEMV at one row, the tensor-core router up to 48.""" + rows = hidden_states.shape[0] if hidden_states.dim() == 2 else 0 + path = dispatch.router_path(rows) + if path == dispatch.ROUTER_GEMV.path: + logits = m3_router_gemv(hidden_states, self.gate.weight, use_pdl=get_env_enable_pdl()) + if logits is not None: + return logits + dispatch.missed(path, rows, "the router GEMV declined the operands") + elif path == dispatch.ROUTER_TC.path: + fused = m3_router_tc(hidden_states, self.gate.weight, use_pdl=get_env_enable_pdl()) + if fused is not None: + return fused[0] + dispatch.missed(path, rows, "the tensor-core router declined the operands") + else: + dispatch.uncovered(dispatch.ROUTER_TC.path, f"more than {ROUTER_TC_MAX_TOKENS} rows") + return self.gate(hidden_states) + + def _router_fp4(self, hidden_states: torch.Tensor): + """Router logits and the routed experts' NVFP4 activation from one kernel, or None (then + ``_router_logits``).""" + if hidden_states.dim() != 2 or hidden_states.shape[0] > ROUTER_TC_MAX_TOKENS: + return None + if self._fp4_input_scale is None: + try: + _, scale = _nvfp4_swiglu_oai_experts(self.experts) + weight = getattr(self.experts, "backend", self.experts).w3_w1_weight + self._fp4_input_scale = ( + scale if weight.shape[-1] * 2 == hidden_states.shape[-1] else False + ) + except ValueError: + self._fp4_input_scale = False + logger.info( + f"[M3 router FP4] fused router+NVFP4 activation quant: {self._fp4_input_scale is not False}" + ) + rows = hidden_states.shape[0] + path = dispatch.router_path(rows) + if path == dispatch.ROUTER_GEMV.path: + path = dispatch.ROUTER_GEMV_FP4.path + if self._fp4_input_scale is False: + dispatch.missed( + path, + rows, + "the routed experts are not TRTLLM-Gen NVFP4 SwiGLU-OAI on this hidden size", + ) + return None + if path == dispatch.ROUTER_TC.path: + fused = m3_router_tc( + hidden_states, self.gate.weight, self._fp4_input_scale, use_pdl=get_env_enable_pdl() + ) + else: + fused = m3_router_gemv_fp4( + hidden_states, self.gate.weight, self._fp4_input_scale, use_pdl=get_env_enable_pdl() + ) + if fused is None: + dispatch.missed(path, rows, "the router declined the operands") + return fused + + def engine(self, num_tokens: int, shared: bool) -> Optional[object]: + """This layer's prebuilt routed-expert engine for ``num_tokens`` (1..8), router folded in. + + ``shared`` folds the shared expert in as well (routed + shared out). None + when the built-in path must run, after a miss: the engine failed to build, + or ``prebuild_engines`` did not build it (a forward never builds one). + """ + key = (num_tokens, True, shared) + return self._lookup(self._engines, key, dispatch.MOE_ENGINE.path, num_tokens) + + def moe16_engine(self, num_tokens: int) -> Optional[object]: + """This layer'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 or self.shared_experts is None: + return None + return self._lookup(self._moe16_engines, num_tokens, dispatch.MOE16_ENGINE.path, num_tokens) + + def _lookup(self, engines: dict, key, path: str, num_tokens: int) -> Optional[object]: + if key not in engines: + what = f"layer {self.layer_idx}'s engine" + if torch.cuda.is_current_stream_capturing(): + dispatch.capture_missed(path, num_tokens, what) + else: + dispatch.missed(path, num_tokens, f"{what} was not built before warmup") + return None + engine = engines[key] + if engine is None: + dispatch.missed(path, num_tokens, "its build failed (logged)") + return engine + + def prebuild_engines(self) -> None: + """Build (compile, on first use of each shared state, and warm) every engine a decode step of this layer can + take: ``MOE_ENGINE``'s token counts with the shared expert folded in, and ``MOE16_ENGINE``'s. Runs from + ``build_targets`` once the weights are real, on the current stream, in the same order on every rank; a build + that fails is logged and its token count falls back (a miss) in the forward. Whether an engine is kept is + agreed across the TP group (one host collective per engine, here): a build that fails on one rank only drops + that engine on every rank, so the ranks never mix an engine (and its boundary push) with the built-in MoE.""" + if self.shared_experts is None: + return + for num_tokens in sorted(dispatch.MOE_ENGINE.keys): + key = (num_tokens, True, True) + if key not in self._engines: + self._engines[key] = self._agreed( + dispatch.MOE_ENGINE.path, + num_tokens, + self._try_build( + dispatch.MOE_ENGINE.path, num_tokens, self._build_engine, shared=True + ), + ) + for num_tokens in sorted(dispatch.MOE16_ENGINE.keys): + if num_tokens not in self._moe16_engines: + self._moe16_engines[num_tokens] = self._agreed( + dispatch.MOE16_ENGINE.path, + num_tokens, + self._try_build( + dispatch.MOE16_ENGINE.path, num_tokens, self._build_moe16_engine + ), + ) + + def _agreed(self, path: str, num_tokens: int, engine: Optional[object]) -> Optional[object]: + return m3_mc_allreduce.tp_agree_built( + self.mapping, + engine, + f"[M3 dispatch] {path} at {num_tokens} rows: layer {self.layer_idx}'s engine", + ) + + def _try_build(self, path: str, num_tokens: int, build, **kwargs) -> Optional[object]: + try: + return build(num_tokens, **kwargs) + except (ImportError, RuntimeError, ValueError) as err: + logger.warning( + f"[M3 dispatch] {path} at {num_tokens} rows: layer {self.layer_idx}'s engine failed to build: {err}" + ) + return None + + def _build_engine(self, num_tokens: int, shared: bool) -> object: + from .ops.moe import m3_moe_engine as m3_moe_op + + tensors, scale = _nvfp4_swiglu_oai_experts(self.experts) + device = tensors[0].device + sms = torch.cuda.get_device_properties(device).multi_processor_count + kwargs = {} + if shared: + gu, dn = self.shared_experts.gate_up_proj, self.shared_experts.down_proj + kwargs["shared_weights"] = ( + gu.weight.data.view(torch.float8_e4m3fn), + gu.weight_scale.data, + dn.weight.data.view(torch.float8_e4m3fn), + dn.weight_scale.data, + ) + state = _ENGINE_STATES.get((str(device), num_tokens, shared)) + if state is None: + state = _ENGINE_STATES[(str(device), num_tokens, shared)] = m3_moe_op.M3MoeState( + device, + num_tokens=num_tokens, + num_ctas=sms - MOE_ENGINE_RESERVED_SMS, + shared=shared, + ) + engine = state.layer( + *tensors, + self.gate.e_score_correction_bias.data, + router_weight=self.gate.weight.data, + fc31_input_scale=scale.reshape(1).float(), + **kwargs, + ) + # Compile (first use of the state) and warm. + engine( + torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device) + ) + return engine + + def _build_moe16_engine(self, num_tokens: int) -> object: + from .ops.moe import m3_moe16_engine as op16 + + tensors, scale = _nvfp4_swiglu_oai_experts(self.experts) + device = tensors[0].device + gu, dn = self.shared_experts.gate_up_proj, self.shared_experts.down_proj + state = _MOE16_STATES.get((str(device), num_tokens)) + if state is None: + state = _MOE16_STATES[(str(device), num_tokens)] = op16.M3Moe16State( + device, num_tokens=num_tokens + ) + engine = state.layer( + *tensors, + self.gate.e_score_correction_bias.data, + router_weight=self.gate.weight.data, + fc31_input_scale=scale.reshape(1).float(), + shared_weights=( + gu.weight.data.view(torch.float8_e4m3fn), + gu.weight_scale.data, + dn.weight.data.view(torch.float8_e4m3fn), + dn.weight_scale.data, + ), + ) + # Compile (first use of the state) and warm. + engine( + torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device) + ) + return engine + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + final_all_reduce_params: Optional[AllReduceParams] = None, + lora_params: Optional[dict] = None, + do_finalize: bool = True, + return_parts: bool = False, + ): + """The MoE output; or, with ``return_parts``, ``(shared, routed)`` (routed None when the + engine folded the shared expert in); or, with ``do_finalize=False``, the unreduced + ``[shared_output, fc2_output, expert_scale_factor, expanded_idx_to_permuted_idx]`` + for the fused finalize + all-reduce + residual + RMSNorm epilogue. + + Mirrors the built-in ``MiniMaxM3MoE.forward`` apart from the decode engines. + """ + all_rank_num_tokens = attn_metadata.all_rank_num_tokens + engine_rows = ( + do_finalize + and hidden_states.dim() == 2 + and hidden_states.dtype == torch.bfloat16 + and 1 <= hidden_states.shape[0] <= MOE_ENGINE_MAX_TOKENS + ) + + def _compute_routed_output(): + fused = self._router_fp4(hidden_states) + if fused is not None: + router_logits, x_fp4, x_sf = fused + return self.experts( + Fp4QuantizedTensor(x_fp4, x_sf, is_sf_swizzled=False), + router_logits, + do_finalize=do_finalize, + output_dtype=hidden_states.dtype, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=False, + lora_params=lora_params, + ) + router_logits = self._router_logits(hidden_states) + return self.experts( + hidden_states, + router_logits, + do_finalize=do_finalize, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=False, + lora_params=lora_params, + ) + + def _compute_shared_output(): + return self.shared_experts(hidden_states, lora_params=lora_params) + + folded = None + if ( + engine_rows + and self.shared_experts is not None + and hidden_states.shape[0] >= MOE_FOLD_MIN_TOKENS + ): + engine = self.engine(hidden_states.shape[0], shared=True) + if engine is not None: + arp = m3_prequant.MOE_PUSH.take_engine() + folded = m3_moe_engine(hidden_states.contiguous(), engine, arp=arp) + if arp is not None: + m3_prequant.MOE_PUSH.engine_pushed = folded.data_ptr() + if folded is not None: + # Routed + shared from one kernel. + if return_parts: + return folded, None + result = folded + elif self.shared_experts is None: + result = _compute_routed_output() + else: + routed_output, shared_output = maybe_execute_in_parallel( + _compute_routed_output, + _compute_shared_output, + self.event_dict[EventType.Main], + self.event_dict[EventType.MoeShared], + # 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 self.aux_stream, + disable_on_compile=True, + ) + if not do_finalize: + return [shared_output, *routed_output] + if return_parts: + return shared_output, routed_output + # In-place add into ``shared_output`` to avoid allocating a temporary. + result = shared_output.add_(routed_output) + # tp4 routed experts are tensor-parallel (never fused-communication), so the sum is still rank-local. + return self.allreduce(result, all_reduce_params=final_all_reduce_params) + + +def producer_pw_covered(layer_uses_nvfp4, num_heads_q, num_heads_kv, num_heads_index) -> bool: + """Whether ``dispatch.PRODUCER`` covers a sparse layer: an NVFP4 KV cache at the TP4 heads (16 query, one KV and + one index head). A covered layer that cannot take ``m3_producer`` is a miss.""" + return bool(layer_uses_nvfp4) and (num_heads_q, num_heads_kv, num_heads_index) == (16, 1, 1) + + +def _producer_miss(covered: bool, rows: int, reason: str) -> None: + """A producer precondition failed: a miss of ``dispatch.PRODUCER`` on a covered layer (raised under + ``modeling_v2='require'``), otherwise nothing; the caller then runs the unfused producer.""" + if covered: + dispatch.missed(dispatch.PRODUCER.path, rows, reason) + return None + + +def producer_pw_supported(num_heads_q, num_heads_kv, num_heads_index, kv_data, kv_scale) -> bool: + """Whether a sparse layer's NVFP4 producer can take ``m3_producer``: the TP4 heads (16 query, one KV and one + index head) and the NVFP4 pool [pages, 2, 1, 128, 64] with its scale pool [pages, 2, 1, 128, 8].""" + return ( + num_heads_q == 16 + and num_heads_kv == 1 + and num_heads_index == 1 + and 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] + ) + + +class M3Attention(MiniMaxM3Attention): + """MiniMax-M3 attention with the decode-step fused kernels.""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + mark_m3_linear(self.qkv_proj, low_m_gemm=True) + mark_m3_linear(self.o_proj, low_m_gemm=True, push=PUSH_ATTN) + self._mega_workspace = None + self._c1_workspace = None + if not self.is_sparse_attention_layer: + # The KV split trtllm-gen dense decode runs at, per request count (msa_decode.dense_decode_max_seq_len). + self.attn.dense_decode_max_seq_len_by_batch = dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH + + def _check_producer_projection(self) -> None: + """A layer the producer covers (dispatch.PRODUCER) needs the fused QKV/index projection: without it the step is + a miss (raises under require) instead of silently taking the unfused path.""" + if not self.enable_fused_qkv_index_projection and producer_pw_covered( + self.main_kv_is_nvfp4, + self.num_heads, + self.num_key_value_heads, + self.sparse_num_index_heads, + ): + dispatch.missed(dispatch.PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _fused_fp8_qkv_indexer_norm_rope_kv_insert( + self, + packed: torch.Tensor, + position_ids: Optional[torch.Tensor], + attn_metadata: AttentionMetadata, + *, + cache_tensors: Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = None, + ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """Run the vLLM-style horizontal producer for every sparse batch. + + Mirrors the built-in ``MiniMaxM3Attention._fused_fp8_qkv_indexer_norm_rope_kv_insert``; an NVFP4 cache takes + this target's ``m3_producer`` (the native producer bit for bit, its constant loads staged before the grid + wait and its dependents triggered on entry) where its geometry is the one that kernel is built for. + + The CUDA kernel is token-major and batch-type agnostic: per-token + positions and cache slots cover pure prefill, mixed aggregate batches, + and CUDA-graph decode. FP8 caches receive FP8 K/V directly. NVFP4 + caches receive packed E2M1 K/V plus their E4M3 SF16 scale bytes while + Q, index-Q, and index-K remain FP8. + """ + if ( + not self.enable_fused_qkv_index_projection + or not isinstance(self.attn, MiniMaxM3MsaSparseAttention) + or not (self.main_kv_is_fp8 or self.main_kv_is_nvfp4) + or self.attn.indexer_kv_dtype != "fp8" + ): + return None + # A sparse layer on an NVFP4 cache at the TP4 heads is the producer's (dispatch.PRODUCER): every precondition + # below that fails for it is a miss. + covered = producer_pw_covered( + self.main_kv_is_nvfp4, + self.num_heads, + self.num_key_value_heads, + self.sparse_num_index_heads, + ) + rows = int(packed.shape[0]) + rope = self.pos_embd_params.rope if self.pos_embd_params is not None else None + rotary_dim = int(rope.dim) if rope is not None else 0 + if ( + packed.dtype != torch.bfloat16 + or position_ids is None + or self.head_dim != 128 + or self.sparse_index_dim != 128 + or self.sparse_num_index_heads != self.num_key_value_heads + or not self.use_gemma_norm + or self.pos_embd_params is None + or not self.pos_embd_params.is_neox + or self.rotary_emb is None + or rope is None + or rotary_dim != 64 + ): + return _producer_miss( + covered, + rows, + "the packed QKV dtype, position ids, head dims, norm or RoPE configuration", + ) + norm_eps = self.q_norm.variance_epsilon + if any( + module.variance_epsilon != norm_eps + for module in (self.k_norm, self.index_q_norm, self.index_k_norm) + ): + return _producer_miss(covered, rows, "the four RMSNorm epsilons differ") + norm_weights = ( + self.q_norm.weight, + self.k_norm.weight, + self.index_q_norm.weight, + self.index_k_norm.weight, + ) + if any(weight.dtype != torch.bfloat16 or not weight.is_cuda for weight in norm_weights): + return _producer_miss(covered, rows, "a norm weight is not BF16 on the GPU") + + kv_cache_manager = attn_metadata.kv_cache_manager + if cache_tensors is None: + if kv_cache_manager is None: + return _producer_miss(covered, rows, "the step has no KV cache manager") + buffers = kv_cache_manager.get_buffers(self.layer_idx, kv_layout="HND") + index_k_cache = attn_metadata.msa_idx_k_cache(self.layer_idx) + out_cache_loc = getattr(attn_metadata, "msa_out_cache_loc", None) + else: + buffers, index_k_cache, out_cache_loc = cache_tensors + num_tokens = int(packed.shape[0]) + layer_uses_nvfp4 = bool( + getattr(kv_cache_manager, "is_nvfp4_layer", lambda _layer_idx: False)(self.layer_idx) + ) + scale_buffers = ( + kv_cache_manager.get_block_scale_buffers(self.layer_idx, kv_layout="HND") + if layer_uses_nvfp4 + else None + ) + kv_quant_scale = getattr(self.qkv_proj, "inv_kv_scales", None) if layer_uses_nvfp4 else None + if layer_uses_nvfp4: + supports_fused_layout = ( + buffers is not None + and scale_buffers is not None + and _nvfp4_cache_supports_fused_write( + buffers, scale_buffers, self.num_key_value_heads + ) + ) + supported_main_cache = ( + supports_fused_layout + and buffers.is_cuda + and scale_buffers.is_cuda + and scale_buffers.device == buffers.device + and kv_quant_scale is not None + 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_cache = ( + buffers is not None + and buffers.is_cuda + and buffers.dtype == torch.float8_e4m3fn + and buffers.dim() == 5 + and tuple(buffers.shape[1:]) == (2, self.num_key_value_heads, 128, 128) + and buffers.stride(4) == 1 + and buffers.stride(3) == 128 + and buffers.stride(2) == 128 * 128 + and buffers.stride(1) >= self.num_key_value_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_cache = ( + 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] + ) + rotary_cos_sin = self.rotary_emb.rotary_cos_sin + supported_rope_cache = ( + rotary_cos_sin.is_cuda + and rotary_cos_sin.dtype == torch.float32 + and rotary_cos_sin.is_contiguous() + and rotary_cos_sin.dim() == 3 + and tuple(rotary_cos_sin.shape[1:]) == (2, 32) + ) + if ( + not supported_main_cache + or not supported_index_cache + or not supported_rope_cache + 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() < num_tokens + ): + return _producer_miss( + covered and layer_uses_nvfp4, + num_tokens, + "the layer's NVFP4 / index caches, RoPE table, KV quant scale or cache slots are not the " + "layout the producer is built for", + ) + + common_args = ( + out_cache_loc[:num_tokens], + self.num_heads, + self.num_key_value_heads, + self.sparse_num_index_heads, + self.head_dim, + rotary_dim, + norm_eps, + self.q_norm.weight, + self.k_norm.weight, + self.index_q_norm.weight, + self.index_k_norm.weight, + rotary_cos_sin, + position_ids.reshape(-1).contiguous().to(torch.int32), + ) + if layer_uses_nvfp4 and producer_pw_supported( + self.num_heads, + self.num_key_value_heads, + self.sparse_num_index_heads, + buffers, + scale_buffers, + ): + q, index_q = m3_producer( + packed.contiguous(), + buffers.view(torch.uint8), + scale_buffers.view(torch.uint8), + index_k_cache, + common_args[0], + kv_quant_scale, + norm_eps, + *common_args[7:12], + common_args[12], + self.num_heads, + get_env_enable_pdl(), + PRODUCER_EARLY_TRIGGER, + ) + elif layer_uses_nvfp4: + if covered: + dispatch.missed( + dispatch.PRODUCER.path, + num_tokens, + f"the NVFP4 pool {tuple(buffers.shape)} / scale pool " + f"{tuple(scale_buffers.shape)} is not the producer's layout", + ) + q, index_q = torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed.contiguous(), + buffers.view(torch.uint8), + scale_buffers, + index_k_cache, + common_args[0], + kv_quant_scale, + *common_args[1:], + ) + else: + q, index_q = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed.contiguous(), + buffers, + index_k_cache, + *common_args, + ) + return q.flatten(1), index_q.flatten(1) + + def _fused_attn_args(self, position_ids, attn_metadata) -> Optional[dict]: + """Arguments of the fused decode attention for this step, or None when it does not apply. + + The layer must be an NVFP4 sparse MSA layer and the step a pure generation + step whose selection the metadata staged as the identity (every request + within 16 blocks of 128 tokens), with at most 8 query tokens per request. + """ + path = dispatch.FUSED_DECODE_ATTN.path + if ( + not isinstance(self.attn, MiniMaxM3MsaSparseAttention) + or not self.is_sparse_attention_layer + or not self.main_kv_is_nvfp4 + or self.attn.indexer_kv_dtype != "fp8" + ): + dispatch.uncovered(path, "not an NVFP4 sparse MSA layer with an FP8 indexer") + return None + if position_ids is None or is_torch_compiling() or attn_metadata.num_contexts > 0: + dispatch.uncovered(path, "not a pure generation step with positions") + return None + if not attn_metadata._msa_short_ok: + dispatch.uncovered( + path, "a request beyond 16 KV blocks: the block selection is not the identity" + ) + return None + span = attn_metadata.msa_decode_span + if span is None or span.row_first != 0: + dispatch.uncovered(path, "no decode span starting at row 0") + return None + num_reqs = int(attn_metadata.num_seqs) + qlen = int(span.query_len) + if not dispatch.FUSED_DECODE_ATTN.covers((qlen, num_reqs)): + dispatch.uncovered( + path, + f"more than {FUSED_ATTN_MAX_QUERY_LEN} query rows per request or " + f"{FUSED_ATTN_MAX_TOKENS} per step", + ) + return None + manager = attn_metadata.kv_cache_manager + if manager is not None and not manager.is_nvfp4_layer(self.layer_idx): + dispatch.uncovered(path, "the layer's KV cache is not NVFP4") + return None + from .kernels import m3_fused_attn + + if m3_fused_attn._module is None: + # Building compiles every geometry (minutes): only outside capture (the eager warmup steps). + if torch.cuda.is_current_stream_capturing(): + dispatch.capture_missed(path, num_reqs * qlen, "the fused attention module") + return None + m3_fused_attn._load_module() + kv_scales = getattr(self.qkv_proj, "kv_scales", None) + inv_kv_scales = getattr(self.qkv_proj, "inv_kv_scales", None) + if manager is None or kv_scales is None or inv_kv_scales is None: + dispatch.missed( + path, + num_reqs * qlen, + "no KV-cache manager" if manager is None else "the QKV projection has no KV scales", + ) + return None + k_scale, v_scale = _aligned_nvfp4_dequant_scales(self.attn, kv_scales) + return dict( + num_tokens=num_reqs * qlen, + decode_query_len=qlen, + kv_data=manager.get_buffers(self.layer_idx, kv_layout="HND").view(torch.uint8), + kv_scale=manager.get_block_scale_buffers(self.layer_idx, kv_layout="HND"), + index_k_cache=attn_metadata.msa_idx_k_cache(self.layer_idx), + out_cache_loc=attn_metadata.msa_out_cache_loc, + position_ids=position_ids.reshape(-1).contiguous().to(torch.int32), + block_table=attn_metadata.msa_block_table[:num_reqs], + seq_lens=attn_metadata.msa_seq_lens_cuda[:num_reqs], + kv_quant_scale=inv_kv_scales, + k_global_scale=k_scale, + v_global_scale=v_scale, + ) + + def _attn_mega(self, hidden_states: torch.Tensor, args: dict) -> Optional[torch.Tensor]: + """o_proj output of the one-launch QKV + attention + o_proj kernel for one request decoding one token whose + layer input the preceding all-reduce handed over in MXFP8; None, with nothing consumed, otherwise.""" + if not dispatch.ATTN_MEGA.covers( + (args["decode_query_len"], int(args["seq_lens"].shape[0])) + ): + return None + if not flashinfer_mxfp8_decode_graph_capture_active(): + return None + pending = m3_prequant.SLOT.pending + if ( + pending is None + or hidden_states.dim() != 2 + or pending[0] != hidden_states.data_ptr() + or pending[1] != 1 + or pending[2] != hidden_states.shape[1] + ): + # After a layer whose boundary does not quantize (a dense layer, an Eagle3 capture layer). + dispatch.uncovered(dispatch.ATTN_MEGA.path, "the layer input has no MXFP8 hand-off") + return None + if self.qkv_proj.weight.shape != (2560, 6144) or self.o_proj.weight.shape != (6144, 2048): + dispatch.uncovered(dispatch.ATTN_MEGA.path, "not the tp4 QKV / o_proj shapes") + return None + from .kernels import m3_attn_mega + + if self._mega_workspace is None: + # The first call JIT-builds the module; both only outside capture (the eager warmup steps). + if torch.cuda.is_current_stream_capturing(): + dispatch.capture_missed( + dispatch.ATTN_MEGA.path, 1, "the attn_mega module and workspace" + ) + return None + m3_attn_mega._load_module() + self._mega_workspace = m3_attn_mega.MegaWorkspace(hidden_states.device) + xq, xsf = m3_prequant.SLOT.take(hidden_states) + ar = m3_prequant.AR_PUSH.take_target() + output = hidden_states.new_empty( + (1, self.num_heads * self.head_dim), dtype=self.attn_activation_dtype + ) + y = hidden_states.new_empty((1, self.o_proj.weight.shape[0]), dtype=torch.bfloat16) + m3_attn_mega.m3_attn_mega( + xq, + xsf, + self.qkv_proj.weight, + self.qkv_proj.weight_scale, + self._mega_workspace, + args["kv_data"], + args["kv_scale"], + args["index_k_cache"], + args["out_cache_loc"], + args["position_ids"], + args["block_table"], + args["seq_lens"], + args["kv_quant_scale"], + args["k_global_scale"], + args["v_global_scale"], + self.q_norm.weight, + self.k_norm.weight, + self.index_k_norm.weight, + self.rotary_emb.rotary_cos_sin, + self.o_proj.weight, + self.o_proj.weight_scale, + eps=self.q_norm.variance_epsilon, + sm_scale=(self.head_dim**-0.5) / float(self.attn.q_scaling), + decode_query_len=1, + output=output, + y=y, + use_pdl=get_env_enable_pdl(), + ar=ar, + ) + if ar is not None: + m3_prequant.AR_PUSH.pushed = y.data_ptr() + return y + + def _attn_c1(self, packed: torch.Tensor, args: dict) -> Optional[torch.Tensor]: + """o_proj output of the one-launch attention + o_proj kernel for one request verifying a draft, pushed into + the attention boundary's all-reduce; None, with nothing consumed, when the step or the layer does not + qualify.""" + qlen, num_tokens = args["decode_query_len"], args["num_tokens"] + if not dispatch.ATTN_C1.covers((qlen, int(args["seq_lens"].shape[0]))): + return None + if not flashinfer_mxfp8_decode_graph_capture_active(): + return None + if m3_prequant.AR_PUSH.target is None: + dispatch.missed( + dispatch.ATTN_C1.path, num_tokens, "no armed attention-boundary push to write into" + ) + return None + if self.o_proj.weight.shape != (6144, 2048): + dispatch.uncovered(dispatch.ATTN_C1.path, "not the tp4 o_proj shape") + return None + from .kernels import m3_attn_mega_c1 + + if self._c1_workspace is None: + # The first call JIT-builds the module; both only outside capture (the eager warmup steps). + if torch.cuda.is_current_stream_capturing(): + dispatch.capture_missed( + dispatch.ATTN_C1.path, num_tokens, "the attn_c1 module and workspace" + ) + return None + m3_attn_mega_c1._load_module() + self._c1_workspace = m3_attn_mega_c1.MegaWorkspace(packed.device) + ar = m3_prequant.AR_PUSH.take_target() + output = packed.new_empty( + (num_tokens, self.num_heads * self.head_dim), dtype=self.attn_activation_dtype + ) + y = packed.new_empty((num_tokens, self.o_proj.weight.shape[0]), dtype=torch.bfloat16) + m3_attn_mega_c1.m3_attn_c1( + packed.contiguous(), + self._c1_workspace, + args["kv_data"], + args["kv_scale"], + args["index_k_cache"], + args["out_cache_loc"], + args["position_ids"], + args["block_table"], + args["seq_lens"], + args["kv_quant_scale"], + args["k_global_scale"], + args["v_global_scale"], + self.q_norm.weight, + self.k_norm.weight, + self.index_k_norm.weight, + self.rotary_emb.rotary_cos_sin, + self.o_proj.weight, + self.o_proj.weight_scale, + eps=self.q_norm.variance_epsilon, + sm_scale=(self.head_dim**-0.5) / float(self.attn.q_scaling), + decode_query_len=qlen, + output=output, + y=y, + use_pdl=get_env_enable_pdl(), + ar=ar[:5], + ) + m3_prequant.AR_PUSH.pushed = y.data_ptr() + return y + + def _fused_attn(self, packed: torch.Tensor, args: dict) -> torch.Tensor: + """The fused norm/RoPE/KV-insert/sparse-decode/merge kernel after the packed QKV GEMM.""" + from .kernels.m3_fused_attn import m3_fused_decode_attn + from .kernels.triton_sparse_decode import MERGE_MXFP8, MERGE_MXFP8_MIN_TOKENS + + num_tokens = args["num_tokens"] + width = self.num_heads * self.head_dim + output = packed.new_empty((num_tokens, width), dtype=self.attn_activation_dtype) + mxfp8 = None + if MERGE_MXFP8_MIN_TOKENS <= num_tokens <= MERGE_MXFP8.MAX_ROWS: + mxfp8 = MERGE_MXFP8.buffers(width, output.device) + m3_fused_decode_attn( + packed.contiguous(), + args["kv_data"], + args["kv_scale"], + args["index_k_cache"], + args["out_cache_loc"], + args["position_ids"], + args["block_table"], + args["seq_lens"], + args["kv_quant_scale"], + args["k_global_scale"], + args["v_global_scale"], + self.q_norm.weight, + self.k_norm.weight, + self.index_k_norm.weight, + self.rotary_emb.rotary_cos_sin, + eps=self.q_norm.variance_epsilon, + sm_scale=(self.head_dim**-0.5) / float(self.attn.q_scaling), + decode_query_len=args["decode_query_len"], + output=output, + mxfp8_out=mxfp8, + use_pdl=get_env_enable_pdl(), + ) + if mxfp8 is not None: + m3_prequant.SLOT.publish(output, mxfp8[0][:num_tokens], mxfp8[1]) + return output + + def _msa_attention_core(self, q, k, v, idx_q, idx_k, attn_metadata, output): + """The built-in MSA attention core, except a pure generation step of a sparse layer whose caches the + fused producer already wrote: that one runs this target's selection and sparse decode.""" + span = attn_metadata.msa_decode_span + if ( + self.is_sparse_attention_layer + and k is None + and attn_metadata.num_contexts == 0 + and span is not None + and span.row_first == 0 + ): + self._decode_sparse_core(q, idx_q, attn_metadata, output) + return output + return super()._msa_attention_core(q, k, v, idx_q, idx_k, attn_metadata, output) + + def _decode_sparse_core(self, q, idx_q, attn_metadata, output) -> None: + """Top-16 block selection and the block-sparse decode of one pure generation step. + + 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). + """ + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + MSA_REQUIRED_TOPK, + msa_paged_kv, + select_blocks_from_maxscore, + ) + + from .kernels.long_select import long_select_blocks + from .kernels.triton_sparse_decode import minimax_m3_sparse_attn_decode + + md = attn_metadata + attn = self.attn + config = attn.m3_config + num_tokens = int(q.shape[0]) + rows = int(md.num_seqs) + qlen = int(md.msa_decode_span.query_len) + block_table = md.msa_block_table[:rows] + seq_lens = md.msa_seq_lens_cuda[:rows] + num_kv_heads = int(attn.indexer.config.num_kv_heads) + # The index scorer's kernels are compiled and loaded by build_targets (prebuild_index_scorer), so no step + # compiles one and the identity can be taken from the first decode step. + identity = md._msa_short_ok and md.msa_short_table is not None + if identity: + topk = md.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 = md.msa_idx_k_cache(self.layer_idx) + max_score = md.msa_proxy_max_score_view( + config.num_index_heads, md.msa_worst_case_max_k_tiles, num_tokens + ) + early = rows <= INDEX_SCORE_EARLY_MAX_REQUESTS + key = _scorer_key( + idx_q_view.dtype, + config.num_index_heads, + qlen, + idx_q_view.shape[-1], + INDEX_SCORE_SPLIT_K, + early, + ) + if key not in _scorer_runner().kernel_cache: + # Outside require mode the op below compiles it (the only scorer this step has). + dispatch.missed( + INDEX_SCORE_PATH, + num_tokens, + f"the index scorer {key} was not compiled before warmup", + ) + 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, + ) + n_valid_blocks = md.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 + ) + 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=num_tokens <= dispatch.SELECT_EARLY_MAX_ROWS, + ) + if table is None: + dispatch.uncovered( + dispatch.LONG_SELECT.path, + f"{max_score_kv.shape[1]} KV blocks, outside {dispatch.LONG_SELECT_MIN_BLOCKS}.." + f"{dispatch.LONG_SELECT_MAX_BLOCKS}: the built-in selector", + ) + table = select_blocks_from_maxscore( + max_score_kv, + topk=MSA_REQUIRED_TOPK, + n_valid_blocks=n_valid_blocks, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + head_major_output=True, + ) + topk = table.permute(1, 0, 2) + manager = md.kv_cache_manager + k_paged, v_paged = msa_paged_kv(manager, self.layer_idx) + nvfp4_args = {} + if manager.is_nvfp4_layer(self.layer_idx): + k_scale, v_scale = _aligned_nvfp4_dequant_scales(attn, self.qkv_proj.kv_scales) + scales = manager.get_block_scale_buffers(self.layer_idx, "HND") + nvfp4_args = dict( + k_block_scale=scales[:, 0], + v_block_scale=scales[:, 1], + k_global_scale=k_scale, + v_global_scale=v_scale, + ) + minimax_m3_sparse_attn_decode( + q.view(num_tokens, self.num_heads, self.head_dim), + k_paged, + v_paged, + topk, + block_table, + seq_lens, + sm_scale=(self.head_dim**-0.5) / float(attn.q_scaling), + output=output.view(num_tokens, self.num_heads, self.head_dim), + decode_query_len=qlen, + identity_blocks=identity, + **nvfp4_args, + ) + + def _dense_decode_fp8_q(self, attn_metadata) -> bool: + """On decode-only steps of a hybrid NVFP4 cache, the dense layers (whose KV pool is FP8) take FP8 q/k/v + straight from the fused norm/RoPE kernel instead of casting them separately.""" + return ( + self.main_kv_is_nvfp4 + and not self.is_sparse_attention_layer + and self._msa_backend_active() + and attn_metadata.num_contexts == 0 + ) + + def _dense_forward( + self, + position_ids: torch.IntTensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + all_reduce_params: Optional[AllReduceParams] = None, + **kwargs, + ) -> torch.Tensor: + """The built-in dense attention (layers 0-2), FP8 q/k/v on decode steps. + + Mirrors the built-in ``MiniMaxM3Attention._dense_forward``. + """ + if attn_metadata is None: + raise RuntimeError( + f"MiniMax-M3 dense forward (layer {self.layer_idx}) requires attn_metadata; received None." + ) + qkv = self.qkv_proj(hidden_states) + fused_qkv = self._fused_qk_norm_rope( + qkv, + position_ids, + num_heads_q=self.num_heads, + num_heads_k=self.num_key_value_heads, + num_heads_v=self.num_key_value_heads, + head_dim=self.head_dim, + q_norm=self.q_norm, + k_norm=self.k_norm, + out_fp8=self._emit_fp8_main_qkv() or self._dense_decode_fp8_q(attn_metadata), + ) + if fused_qkv is not None: + q, k, v = self._split_main_qkv(fused_qkv) + else: + assert not self._expect_fused_qk_norm_rope(position_ids), ( + f"MiniMax-M3 dense attention (layer {self.layer_idx}) expected the " + f"fused QK-norm+RoPE kernel (bf16 activations, head_dim=" + f"{self.head_dim}) but fell back to the separate path; qkv dtype " + f"is {qkv.dtype} (expected {self.attn_activation_dtype})." + ) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.apply_qk_norm(q, k) + if self.rotary_emb is not None and position_ids is not None: + q, k = self.rotary_emb(position_ids, [q, k]) + o = self._forward_attention_core(q, k, v, None, None, attn_metadata) + return self.o_proj(o, all_reduce_params=all_reduce_params) + + def _sparse_forward( + self, + position_ids: torch.IntTensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + all_reduce_params: Optional[AllReduceParams] = None, + **kwargs, + ) -> torch.Tensor: + """The built-in sparse attention, with the fused decode kernels on the steps they take. + + Mirrors the built-in ``MiniMaxM3Attention._sparse_forward``. + """ + if attn_metadata is None: + raise RuntimeError( + f"MiniMax-M3 sparse forward (layer {self.layer_idx}) requires attn_metadata; received None." + ) + packed_qkv = None + packed_idx_qk = None + self._check_producer_projection() + if self.enable_fused_qkv_index_projection: + if ( + self.register_to_config + and is_torch_compiling() + and isinstance(self.attn, MiniMaxM3MsaSparseAttention) + and self._emit_fp8_main_qkv() + and self.attn.indexer_kv_dtype == "fp8" + ): + kv_cache, index_k_cache = attn_metadata.msa_layer_cache_tensors[self.layer_idx] + q, idx_q = torch.ops.trtllm.minimax_m3_fused_sparse_qkv_producer( + hidden_states, + position_ids, + kv_cache, + index_k_cache, + attn_metadata.msa_out_cache_loc, + self.layer_idx_str, + ) + o = self._forward_attention_core(q, None, None, idx_q, None, attn_metadata) + return self.o_proj(o, all_reduce_params=all_reduce_params) + if self.register_to_config and (is_torch_compiling() or is_in_breakable_cuda_graph()): + packed = torch.ops.trtllm.minimax_m3_qkv_index_proj( + hidden_states, position_ids, self.layer_idx_str + ) + num_tokens = ( + position_ids.shape[-1] if position_ids is not None else hidden_states.shape[0] + ) + output = packed.new_empty( + (num_tokens, self.num_heads * self.head_dim), dtype=self.attn_activation_dtype + ) + maybe_bcg_minimax_m3_attn_custom_op_inplace( + None, None, None, None, None, packed, position_ids, self.layer_idx_str, output + ) + return self.o_proj(output, all_reduce_params=all_reduce_params) + # Gathered before the QKV GEMM so nothing runs between it and the fused attention (PDL hand-off). + fused_attn = self._fused_attn_args(position_ids, attn_metadata) + if fused_attn is not None: + y = self._attn_mega(hidden_states, fused_attn) + if y is not None: + return y + packed = self.qkv_proj(hidden_states) + if fused_attn is not None and ( + packed.shape[0] != fused_attn["num_tokens"] or packed.shape[1] != 2560 + ): + dispatch.missed( + dispatch.FUSED_DECODE_ATTN.path, + fused_attn["num_tokens"], + f"the packed QKV is {tuple(packed.shape)}, not [rows, 2560]", + ) + fused_attn = None + if fused_attn is not None: + y = self._attn_c1(packed, fused_attn) + if y is not None: + return y + o = self._fused_attn(packed, fused_attn) + return self.o_proj(o, all_reduce_params=all_reduce_params) + horizontal = self._fused_fp8_qkv_indexer_norm_rope_kv_insert( + packed, position_ids, attn_metadata + ) + if horizontal is not None: + q, idx_q = horizontal + o = self._forward_attention_core(q, None, None, idx_q, None, attn_metadata) + return self.o_proj(o, all_reduce_params=all_reduce_params) + main_size = self.q_size + 2 * self.kv_size + packed_qkv, packed_idx_qk = packed.split( + [main_size, self.index_q_size + self.index_k_size], dim=-1 + ) + + # The compatibility/fallback main and index branches, overlapped on the aux stream. + def _main_norm_rope(): + qkv = packed_qkv if packed_qkv is not None else self.qkv_proj(hidden_states) + fused_qkv = self._fused_qk_norm_rope( + qkv, + position_ids, + num_heads_q=self.num_heads, + num_heads_k=self.num_key_value_heads, + num_heads_v=self.num_key_value_heads, + head_dim=self.head_dim, + q_norm=self.q_norm, + k_norm=self.k_norm, + out_fp8=self._emit_fp8_main_qkv(), + ) + if fused_qkv is not None: + return self._split_main_qkv(fused_qkv) + assert not self._expect_fused_qk_norm_rope(position_ids), ( + f"MiniMax-M3 sparse attention (layer {self.layer_idx}) expected the fused QK-norm+RoPE kernel" + ) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.apply_qk_norm(q, k) + if self.rotary_emb is not None and position_ids is not None: + q, k = self.rotary_emb(position_ids, [q, k]) + return q, k, v + + def _index_norm_rope(): + idx_qk = ( + packed_idx_qk if packed_idx_qk is not None else self.index_qk_proj(hidden_states) + ) + graph_fp8_indexer = ( + self.register_to_config + and is_torch_compiling() + and isinstance(self.attn, MiniMaxM3MsaSparseAttention) + and self.attn.indexer_kv_dtype == "fp8" + ) + if not graph_fp8_indexer: + fp8_idx_q = self._fused_fp8_index_qk_norm_rope(idx_qk, position_ids, attn_metadata) + if fp8_idx_q is not None: + return fp8_idx_q, None + fused_idx = self._fused_qk_norm_rope( + idx_qk, + position_ids, + num_heads_q=self.sparse_num_index_heads, + num_heads_k=1, + num_heads_v=0, + head_dim=self.sparse_index_dim, + q_norm=self.index_q_norm, + k_norm=self.index_k_norm, + out_fp8=graph_fp8_indexer, + ) + if fused_idx is not None: + return self._split_index_qk(fused_idx) + assert not self._expect_fused_qk_norm_rope(position_ids), ( + f"MiniMax-M3 sparse index branch (layer {self.layer_idx}) expected the fused QK-norm+RoPE kernel" + ) + idx_q, idx_k = idx_qk.split([self.index_q_size, self.index_k_size], dim=-1) + idx_q, idx_k = self.apply_index_qk_norm(idx_q, idx_k) + if self.rotary_emb is not None and position_ids is not None: + idx_q, idx_k = self.rotary_emb(position_ids, [idx_q, idx_k]) + return idx_q, idx_k + + (q, k, v), (idx_q, idx_k) = maybe_execute_in_parallel( + _main_norm_rope, + _index_norm_rope, + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + disable_on_compile=True, + ) + o = self._forward_attention_core(q, k, v, idx_q, idx_k, attn_metadata) + return self.o_proj(o, all_reduce_params=all_reduce_params) + + +def _plain_capture(spec_metadata, layer_id: int, hidden_states, residual) -> None: + """The in-tree Eagle3 hidden-state capture (no fused fc_norm, no early fc fork).""" + capture_hidden_states(spec_metadata, layer_id, hidden_states, residual, None) + + +class M3DecoderLayer(DecoderLayer): + """One MiniMax-M3 transformer block on this target's modules. + + Mirrors the built-in ``MiniMaxM3DecoderLayer``; the layer boundaries add the + multicast all-reduce + residual + RMSNorm up to its token limits. + """ + + def __init__( + self, + model_config: ModelConfig[PretrainedConfig], + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ) -> None: + super().__init__() + config = model_config.pretrained_config + self.hidden_size = config.hidden_size + self.layer_idx = layer_idx + self.mapping = model_config.mapping + self.enable_attention_dp = self.mapping.enable_attention_dp + + _, sparse_layer_ids = get_sparse_layer_ids(config) + disable_index_value_ids = set(get_sparse_disable_index_value_layer_ids(config)) + self.self_attn = M3Attention( + model_config=model_config, + layer_idx=layer_idx, + is_sparse_attention_layer=layer_idx in sparse_layer_ids, + disable_index_value=layer_idx in disable_index_value_ids, + aux_stream=aux_stream_dict[AuxStreamType.Attention], + ) + + _, moe_layer_ids = get_moe_layer_ids(config) + if layer_idx in moe_layer_ids: + self.block_sparse_moe = M3MoE( + model_config=model_config, aux_stream_dict=aux_stream_dict, layer_idx=layer_idx + ) + self.mlp = None + else: + dense_intermediate = int( + getattr(config, "dense_intermediate_size", config.intermediate_size) + ) + self.mlp = _dense_mlp(model_config, dense_intermediate, layer_idx) + self.block_sparse_moe = None + + # Plain (non-Gemma) boundary norms, so they can drive the fused all-reduce + # + residual + RMSNorm epilogue; the loader folds Gemma's (1 + weight) in. + self.input_layernorm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + use_gemma=False, + ) + self.post_attention_layernorm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + use_gemma=False, + ) + + # tp4 without attention DP: every boundary has a cross-rank reduction to fold. + self.pre_feed_forward_fusion = True + self.post_feed_forward_fusion = True + self.allreduce = AllReduce( + mapping=model_config.mapping, + strategy=model_config.allreduce_strategy, + dtype=config.torch_dtype, + ) + # Wired after the weight load to the next layer's input_layernorm (or the final norm). + self.next_layer_layernorm: Optional[RMSNorm] = None + self.next_is_decoder_layer = False + # The Eagle3 hidden-state capture of this layer; the target shell points it at its worker's (fused fc_norm, + # early fc fork). + self.capture: Callable[..., None] = _plain_capture + + def forward( + self, + position_ids: torch.IntTensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + residual: Optional[torch.Tensor], + spec_metadata: Optional[SpecMetadata] = None, + lora_params: Optional[dict] = None, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + all_reduce_params=AllReduceParams(enable_allreduce=False), + lora_params=lora_params, + **kwargs, + ) + capture = spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx) + if self.block_sparse_moe is not None: + return self.forward_moe( + hidden_states, attn_metadata, residual, capture, spec_metadata, lora_params + ) + return self.forward_mlp(hidden_states, residual, capture, spec_metadata, lora_params) + + def _can_mc(self, hidden_states: torch.Tensor, folded: bool = False) -> bool: + """Whether the multicast all-reduce + residual + RMSNorm takes this boundary.""" + return ( + isinstance(hidden_states, torch.Tensor) + and hidden_states.dim() == 2 + and hidden_states.shape[0] + <= (m3_mc_allreduce.MAX_TOKENS if folded else MC_AR_MAX_TOKENS) + and hidden_states.dtype == torch.bfloat16 + and hidden_states.is_contiguous() + ) + + def mc_allreduce_norm(self, x, x2, residual, norm): + """The multicast all-reduce + optional second addend + residual + RMSNorm, or None without a workspace.""" + workspace = m3_mc_allreduce.get_workspace(self.mapping, x.shape[-1]) + if workspace is None: + return None + return m3_mc_allreduce.m3_mc_allreduce_norm( + x, + residual.contiguous(), + norm.weight, + norm.variance_epsilon, + workspace, + use_pdl=get_env_enable_pdl(), + x2=x2, + ) + + def native_allreduce_norm(self, hidden_states, residual, norm): + """The built-in fused all-reduce + residual + RMSNorm.""" + return self.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=norm.weight, + eps=norm.variance_epsilon, + trigger_completion_at_end=False, + ), + ) + + def _pre_feed_forward_norm(self, hidden_states, residual): + if self._can_mc(hidden_states): + fused = self.mc_allreduce_norm( + hidden_states, None, residual, self.post_attention_layernorm + ) + if fused is not None: + return fused + return self.native_allreduce_norm(hidden_states, residual, self.post_attention_layernorm) + + def _next_layer_norm(self, hidden_states, residual, fused_boundary: bool): + """The next layer's input norm at the boundary; ``fused_boundary`` False keeps the module's own reduction.""" + if self.next_layer_layernorm is None: + return hidden_states, residual + parts = hidden_states if isinstance(hidden_states, tuple) else (hidden_states, None) + if not fused_boundary: + return self.next_layer_layernorm(parts[0], residual) + if self._can_mc(parts[0], folded=parts[1] is not None): + fused = self.mc_allreduce_norm(parts[0], parts[1], residual, self.next_layer_layernorm) + if fused is not None: + return fused + hidden_states = parts[0] if parts[1] is None else parts[0].add_(parts[1]) + return self.native_allreduce_norm(hidden_states, residual, self.next_layer_layernorm) + + def _fused_feed_forward_boundary(self, capture: bool) -> bool: + """Whether the feed-forward output goes unreduced to the fused boundary (all-reduce + residual + next norm in + one kernel). An Eagle3 capture layer takes it too and captures the boundary's reduced residual sum; only the + last layer (no next norm, its raw output captured) keeps the module's own all-reduce.""" + return not capture or self.next_layer_layernorm is not None + + def _feed_forward_boundary( + self, hidden_states, residual, fused_boundary, capture, spec_metadata + ): + """The layer's exit: the next norm at the boundary, with this layer's Eagle3 capture.""" + if not capture: + return self._next_layer_norm(hidden_states, residual, fused_boundary) + if not fused_boundary: + self.capture(spec_metadata, self.layer_idx, hidden_states, residual) + return self._next_layer_norm(hidden_states, residual, False) + normed, updated = self._next_layer_norm(hidden_states, residual, True) + self.capture(spec_metadata, self.layer_idx, updated, None) + return normed, updated + + def forward_moe( + self, hidden_states, attn_metadata, residual, capture, spec_metadata, lora_params + ): + hidden_states, residual = self._pre_feed_forward_norm(hidden_states, residual) + fused_boundary = self._fused_feed_forward_boundary(capture) + hidden_states = self.block_sparse_moe( + hidden_states, + attn_metadata, + final_all_reduce_params=AllReduceParams(enable_allreduce=False) + if fused_boundary + else None, + lora_params=lora_params, + return_parts=fused_boundary and self._can_mc(hidden_states, folded=True), + ) + return self._feed_forward_boundary( + hidden_states, residual, fused_boundary, capture, spec_metadata + ) + + def forward_mlp(self, hidden_states, residual, capture, spec_metadata, lora_params): + hidden_states, residual = self._pre_feed_forward_norm(hidden_states, residual) + fused_boundary = self._fused_feed_forward_boundary(capture) + hidden_states = self.mlp( + hidden_states, + final_all_reduce_params=AllReduceParams(enable_allreduce=False) + if fused_boundary + else None, + lora_params=lora_params, + ) + return self._feed_forward_boundary( + hidden_states, residual, fused_boundary, capture, spec_metadata + ) + + +class M3TextModel(DecoderModel): + """The MiniMax-M3 text decoder on this target's layers. + + Its forward belongs to the core in ``modeling.py``, which dispatches each + step to the prefill or the decode target. + """ + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + _validate_sparse_attention_runtime_config(model_config) + super().__init__(model_config) + config = model_config.pretrained_config + self.vocab_size = config.vocab_size + self.aux_stream_dict = { + AuxStreamType.Attention: torch.cuda.Stream(), + AuxStreamType.MoeShared: torch.cuda.Stream(), + AuxStreamType.MoeChunkingOverlap: torch.cuda.Stream(), + } + self.embed_tokens = Embedding( + config.vocab_size, + config.hidden_size, + dtype=config.torch_dtype, + enable_torch_compile_for_embedding=model_config.enable_torch_compile_for_embedding, + ) + self.layers = nn.ModuleList( + [ + M3DecoderLayer(model_config, layer_idx, self.aux_stream_dict) + for layer_idx in range(config.num_hidden_layers) + ] + ) + # Lets the executor enable the automatic FlashInfer MXFP8 path for the decode graphs. + self._use_flashinfer_mxfp8_decode_graph_default = True + self.norm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + use_gemma=False, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/linears.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/linears.py new file mode 100644 index 000000000000..a36255fa92e2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/linears.py @@ -0,0 +1,259 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-graph MXFP8 GEMM dispatch of this target's projections. + +Every MXFP8 projection of the text model (QKV, o_proj, the dense MLPs and the +shared experts) runs the built-in ``MXFP8LinearMethod`` outside decode CUDA +graphs. Inside them (``flashinfer_mxfp8_decode_graph_capture_active``) the +projection's (N, K) and the step's rows pick one path (``dispatch.linear_path``, +row ranges per projection in ``dispatch.py``): + +1. an activation its producer already quantized to MXFP8 (the boundary + all-reduce or the attention merge publish it in ``m3_prequant.SLOT``): the + tcgen05 decode GEMV on the rows it serves, else the decode-graph MXFP8 GEMM + (AutoTuner-selected backend per token bucket, the built-in method's GEMM); +2. a bf16 activation on the rows the GEMV serves quantizing in its prologue; +3. projections marked for it, on its rows: the fused low-M MXFP8 GEMM; +4. otherwise the built-in method (MXFP8 quantize + GEMM, AutoTuner-selected per bucket). + +No path runs FlashInfer's own per-rank config sweep: its tactics are not merged +across ranks and it profiled wherever the projection ran (the MoE aux stream for +the shared expert), one CUDA graph per tactic. + +A kernel declining rows its path covers is a miss (``dispatch.missed``); the +chain then continues down the list. The GEMVs push their rows straight into +the boundary all-reduce the decode forward armed for the projection's push +slot (``m3_push``, fixed at construction). +""" + +from __future__ import annotations + +from typing import List, Optional, Tuple + +import torch +from torch import nn + +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_gemm_autotuned +from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.mx_gemv import op as mx_gemv +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.modules.linear import ( + Linear, + MXFP8LinearMethod, + flashinfer_mxfp8_autotune_active, + flashinfer_mxfp8_decode_graph_capture_active, +) +from tensorrt_llm._torch.utils import is_torch_compiling + +from . import dispatch +from .kernels import m3_mc_allreduce, m3_prequant +from .kernels.mxfp8_low_m_gemm import mxfp8_low_m_gemm +from .ops.gemm.m3_mx_gemv import m3_mx_gemv +from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 + +# The boundary all-reduce a projection's GEMV pushes its rows into, when the decode forward arms one: the attention +# boundary (o_proj) or the MoE boundary's shared-expert set (the shared expert's down projection). +PUSH_ATTN = "attn" +PUSH_MOE = "moe" + + +def mark_m3_linear(module: Linear, low_m_gemm: bool, push: Optional[str] = None) -> None: + """Mark a projection for this target's decode dispatch; ``adopt_m3_linears`` applies it after the load. + + ``low_m_gemm``: the projection may take the fused low-M MXFP8 GEMM. ``push``: ``PUSH_ATTN`` / ``PUSH_MOE`` when + its GEMV pushes into that boundary's all-reduce whenever the decode forward armed it. + """ + module.m3_low_m_gemm = low_m_gemm + module.m3_push = push + + +def adopt_m3_linears(root: nn.Module) -> None: + """Put every marked projection under ``root`` on the decode dispatch. + + Runs once the weights are real: the quantization method a projection ends + up with is only final after the model applied its quantization config. + It also compiles every decode GEMV kernel a projection can reach here, + before the executor's warmup: on the current (main) stream, the same shapes + in the same order on every rank, and never inside a forward between + collectives or on the MoE aux stream (``prebuild_m3_gemvs``). + """ + use_pdl = get_env_enable_pdl() + compiled = False + for name, module in root.named_modules(): + if not hasattr(module, "m3_low_m_gemm"): + continue + method = module.quant_method + if not isinstance(method, MXFP8LinearMethod) or not method.use_cutlass: + raise ValueError( + f"{name}: the MiniMax-M3 decode dispatch needs a CUTLASS MXFP8 projection, got {type(method).__name__}" + ) + if type(method) is MXFP8LinearMethod: + method.__class__ = M3MXFP8LinearMethod + mx_gemv.ensure_compiled(module.weight, module.weight_scale, use_pdl) + prebuild_m3_gemvs(module, use_pdl) + compiled = True + if compiled: + torch.cuda.synchronize() + + +def _bucket_tokens(rows: int) -> int: + """The token capacity of the multicast all-reduce bucket a ``rows``-token boundary uses.""" + return next(b for b in m3_mc_allreduce.BUCKETS if rows <= b) + + +def gemv_variants(module: Linear) -> List[Tuple[int, bool, Optional[Tuple[int, int, int]]]]: + """Every (rows, bf16 activation, push (world, bucket tokens, mode)) call the decode dispatch can route to this + projection's GEMV (``M3MXFP8LinearMethod.apply``): 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.""" + n, k = module.weight.shape + push_rows = {PUSH_ATTN: dispatch.ATTN_PUSH.keys, PUSH_MOE: dispatch.MOE_PUSH.keys}.get( + module.m3_push, frozenset() + ) + world = int(module.tp_size) + variants = [] + for rows in range(1, dispatch.DECODE_MAX_TOKENS + 1): + for x_bf16, served in ( + (False, rows in dispatch.MX_GEMV_ROWS.get((n, k), ())), + (True, rows <= dispatch.MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0)), + ): + if not served: + continue + variants.append((rows, x_bf16, None)) + if rows in push_rows: + variants.append( + (rows, x_bf16, (world, _bucket_tokens(rows), dispatch.AR_PUSH_MODE)) + ) + return variants + + +def prebuild_m3_gemvs(module: Linear, use_pdl: bool) -> int: + """Compile every GEMV kernel of ``gemv_variants`` (the push ones included) without launching; returns how many + distinct kernels the projection takes.""" + keys = set() + for rows, x_bf16, push in gemv_variants(module): + key = mx_gemv.precompile( + module.weight, module.weight_scale, rows, x_bf16, use_pdl=use_pdl, ar=push + ) + if key is not None: + keys.add(key) + return len(keys) + + +def _low_m_gemm_supports(x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor) -> bool: + """Whether the fused MXFP8 low-M kernel handles these operands.""" + return ( + x.dim() == 2 + and x.dtype in (torch.bfloat16, torch.float16) + and x.stride(-1) == 1 + and weight.dtype == torch.float8_e4m3fn + and weight.is_contiguous() + and weight_scale.dtype == torch.uint8 + and x.shape[1] == weight.shape[1] + and weight.shape[1] % 256 == 0 + ) + + +def _push_target(module: Linear): + """The armed boundary all-reduce this projection pushes its output into, if any.""" + if module.m3_push == PUSH_MOE: + return m3_prequant.MOE_PUSH.take_shared() + if module.m3_push == PUSH_ATTN: + return m3_prequant.AR_PUSH.take_target() + return None + + +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 _note_pushed(output: Optional[torch.Tensor], module: Linear) -> None: + if output is None or not mx_gemv.last_pushed: + return + if module.m3_push == PUSH_MOE: + m3_prequant.MOE_PUSH.shared_pushed = output.data_ptr() + else: + m3_prequant.AR_PUSH.pushed = output.data_ptr() + + +class M3MXFP8LinearMethod(MXFP8LinearMethod): + """``MXFP8LinearMethod`` with the decode-graph dispatch in the module docstring.""" + + def apply(self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor]): + if bias is not None or input.dim() != 2: + return super().apply(module, input, bias) + input = input.contiguous() + weight, weight_scale = module.weight, module.weight_scale + use_pdl = get_env_enable_pdl() + decode_graph = not is_torch_compiling() and flashinfer_mxfp8_decode_graph_capture_active() + if not decode_graph: + return super().apply(module, input, bias) + + rows = input.shape[0] + n, k = weight.shape + prequant = m3_prequant.SLOT.take(input) + if module.dtype != torch.bfloat16 or (prequant is None and input.dtype != torch.bfloat16): + path = ( + dispatch.PREQUANT_GEMM_PATH if prequant is not None else dispatch.GENERIC_GEMM_PATH + ) + else: + path = dispatch.linear_path(n, k, rows, prequant is not None, module.m3_low_m_gemm) + output = None + if prequant is not None: + act, act_scale = prequant + if path == dispatch.MX_GEMV_PATH: + output = m3_mx_gemv( + act, act_scale, weight, weight_scale, use_pdl=use_pdl, ar=_push_target(module) + ) + _note_pushed(output, module) + if output is None: + dispatch.missed( + path, + rows, + f"the GEMV declined the [{n}, {k}] projection's MXFP8 activation" + + _uncompiled_note(), + ) + if output is None: + if path == dispatch.PREQUANT_GEMM_PATH: + dispatch.uncovered( + dispatch.MX_GEMV_PATH, + f"the [{n}, {k}] projection at rows the GEMV does not serve", + ) + output = mxfp8_gemm_autotuned( + act, + act_scale, + weight, + weight_scale, + module.dtype, + tune=flashinfer_mxfp8_autotune_active(), + ) + return output + if path == dispatch.MX_GEMV_BF16_PATH: + output = m3_mx_gemv_bf16( + input, weight, weight_scale, use_pdl=use_pdl, ar=_push_target(module) + ) + _note_pushed(output, module) + if output is None: + dispatch.missed( + path, + rows, + f"the GEMV declined the [{n}, {k}] projection's input" + _uncompiled_note(), + ) + if ( + output is None + and path in (dispatch.MX_GEMV_BF16_PATH, dispatch.LOW_M_GEMM_PATH) + and module.m3_low_m_gemm + and rows <= dispatch.LOW_M_GEMM_MAX_TOKENS.get((n, k), 0) + and _low_m_gemm_supports(input, weight, weight_scale) + ): + output = mxfp8_low_m_gemm(input, weight, weight_scale, module.dtype, use_pdl=use_pdl) + if output is None and path == dispatch.LOW_M_GEMM_PATH: + dispatch.missed(path, rows, f"the GEMM declined the [{n}, {k}] projection's input") + if output is None: + if path == dispatch.GENERIC_GEMM_PATH: + dispatch.uncovered( + dispatch.MX_GEMV_BF16_PATH, + f"the [{n}, {k}] projection's bf16 input beyond the decode GEMV's rows", + ) + # Stock's decode-graph dispatch for this shape (mxfp8_quantize_gemm_autotuned): same tactics, same bits. + output = super().apply(module, input, bias) + return output diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py new file mode 100644 index 000000000000..219d0a1e0300 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py @@ -0,0 +1,633 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ModelingV2 target: MiniMax-M3 NVFP4 / sm_103 / tp4. + +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) +plus one shared expert. The checkpoint is a modelopt mixed-precision export: +NVFP4 routed experts, MXFP8 projections, bf16 router, embedding and head; the KV +cache is NVFP4 for the sparse layers and FP8 for the dense ones. + +``tp4`` is four GB300s, head-split attention (16 query heads and 1 KV head per +rank) and tensor-parallel experts (every expert on every rank, width 768 per +rank), no attention DP. Speculative decoding, when configured, is one-model +Eagle3 with draft length 3: the in-tree BF16 drafter over the full vocabulary +and the BF16 target head, both at the checkpoint's precision; decode steps of +up to 8 requests verify and draft through fused kernels at that precision +(``spec_worker.py``). + +**Each step is dispatched once, on the host, from its shape** (``_select_target``): + +* ``DecodeTarget`` takes a pure generation step of at most 48 tokens (up to 8 + requests, each one token or a draft verify of 1 + 3 + extra rows). It runs + the layers with the decode-step kernels of ``layers.py`` and chooses each + layer boundary's all-reduce by the step's token bucket (``DECODE_PLANS``). + Every fast path it can take, the rows each covers and what a step that cannot + take one does (raise under ``modeling_v2='require'``, else log once and fall + back) are in ``dispatch.py``. +* ``PrefillTarget`` takes every other step (prefill, mixed steps, wider decode + steps) on the built-in modules' generic path. + +**What this target asserts rather than adapts**: SM 10.3, the topology above, +the checkpoint shape and quantization (``routing.py``), the MiniMax-M3 MSA +sparse attention backend with its NVFP4 hybrid KV cache, and text-only input: +the vision tower is not built and a step carrying image input raises. +""" + +from __future__ import annotations + +import time +from typing import Any, Dict, List, Optional + +import torch +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + ModelingV2Mode, + resolved_mode, +) +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.mx_gemv import op as mx_gemv +from tensorrt_llm._torch.distributed import 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_minimaxm3 import ( + MiniMaxM3ForCausalLM, + get_text_model_config, + is_minimax_m3_vl_config, +) +from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM +from tensorrt_llm._torch.modules.linear import flashinfer_mxfp8_decode_graph_capture_active +from tensorrt_llm._torch.speculative import SpecMetadata +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantAlgo + +from . import dispatch +from . import weights as _weights +from .dispatch import AR_PUSH_MODE, EMBED_FIRSTNORM_MAX_TOKENS, DecodePlan, plan_for +from .kernels import m3_mc_allreduce, m3_prequant +from .layers import M3DecoderLayer, M3MoE, M3TextModel, prebuild_index_scorer +from .linears import adopt_m3_linears +from .ops.moe.m3_moe16_engine import m3_moe16_engine +from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm +from .runtime import make_runtime_hooks +from .spec_worker import M3Eagle3Worker, M3SpecDecForwardMixin + +_SM = (10, 3) + + +class MiniMaxM3Tp4ModelingV2(M3TextModel): + """The text model of this target: owns the layers and dispatches each step to a target.""" + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + super().__init__(model_config) + self._embed_owners: Dict[tuple, DecodeEmbedding] = {} + self._embed_identity: Optional[tuple] = None + self._prefill: Optional[PrefillTarget] = None + self._decode: Optional[DecodeTarget] = None + + def build_targets(self) -> None: + """Build the two targets, once the weights are real.""" + table = self.embed_tokens.weight + norm = self.layers[0].input_layernorm + identity = ( + id(table), + table.data_ptr(), + id(norm.weight), + norm.weight.data_ptr(), + float(norm.variance_epsilon), + ) + if self._embed_owners and identity != self._embed_identity: + raise RuntimeError( + "the embedding or the first norm moved after the decode embedding owners were built" + ) + if not self._embed_owners: + self._embed_owners = { + (rows, dtype): DecodeEmbedding(table, norm, rows, dtype) + for rows in range(1, EMBED_FIRSTNORM_MAX_TOKENS + 1) + for dtype in (torch.int32, torch.int64) + } + self._embed_identity = identity + self.prebuild_kernels() + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + + 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 ``dispatch.missed``). + Returns the seconds each part took and the scorer keys.""" + times: Dict[str, object] = {} + mx_gemv.freeze(False) + start = time.perf_counter() + adopt_m3_linears(self) + times["gemv_s"] = time.perf_counter() - start + start = time.perf_counter() + for layer in self.layers: + moe = layer.block_sparse_moe + if isinstance(moe, M3MoE): + moe.prebuild_engines() + torch.cuda.synchronize() + times["moe_s"] = time.perf_counter() - start + start = time.perf_counter() + times["scorer_keys"] = prebuild_index_scorer( + self.layers, self.decode_query_lens(), self.embed_tokens.weight.device + ) + torch.cuda.synchronize() + times["scorer_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)" + % (times["gemv_s"], times["moe_s"], times["scorer_s"], len(times["scorer_keys"])) + ) + return times + + def _select_target(self, attn_metadata: AttentionMetadata, input_ids: Optional[torch.Tensor]): + """Which target runs this step: host-side counts only, fixed per CUDA-graph capture.""" + if ( + attn_metadata.num_contexts == 0 + and input_ids is not None + and input_ids.ndim == 1 + and dispatch.is_decode_step(input_ids.numel()) + ): + return self._decode + return 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, + **kwargs, + ) -> torch.Tensor: + 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") + target = ( + self._select_target(attn_metadata, input_ids) + if inputs_embeds is None + else self._prefill + ) + return target.forward( + attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, **kwargs + ) + + +class PrefillTarget: + """The general case: context rows, mixed steps and wide generation steps, on the generic layer path.""" + + def __init__(self, core: MiniMaxM3Tp4ModelingV2) -> None: + self.core = core + + def forward( + self, attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, **kwargs + ): + core = self.core + hidden_states = core.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds + residual = None + for layer in core.layers: + hidden_states, residual = layer( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + residual=residual, + spec_metadata=spec_metadata, + **kwargs, + ) + return hidden_states + + +class DecodeTarget: + """A pure generation step of at most 48 tokens. + + 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. ``dispatch.py`` lists the row + ranges and what a covered step that cannot take its path does. + """ + + def __init__(self, core: MiniMaxM3Tp4ModelingV2) -> None: + self.core = core + self._moe_push_descs: Dict[tuple, Any] = {} + + def forward( + self, attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, **kwargs + ): + core = self.core + num_tokens = int(input_ids.numel()) + plan = plan_for(num_tokens) + hidden_states, residual = self._embed(input_ids, num_tokens) + for layer in core.layers: + if residual is None or ( + spec_metadata is not None and spec_metadata.is_layer_capture(layer.layer_idx) + ): + hidden_states, residual = layer( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + residual=residual, + spec_metadata=spec_metadata, + **kwargs, + ) + else: + hidden_states, residual = self._layer( + layer, plan, position_ids, hidden_states, residual, attn_metadata, kwargs + ) + return hidden_states + + 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 dispatch.EMBED_FIRSTNORM.covers(num_tokens): + dispatch.uncovered( + dispatch.EMBED_FIRSTNORM.path, f"more than {EMBED_FIRSTNORM_MAX_TOKENS} tokens" + ) + return core.embed_tokens(input_ids), None + owner = core._embed_owners.get((num_tokens, input_ids.dtype)) + if owner is None or not input_ids.is_contiguous(): + dispatch.missed( + dispatch.EMBED_FIRSTNORM.path, + num_tokens, + f"no embedding owner for {input_ids.dtype} ids (built for int32 / int64 by build_targets)" + if owner is None + else "the input ids are not contiguous", + ) + return core.embed_tokens(input_ids), None + return m3_embedding_firstnorm( + input_ids, + owner.table, + owner.norm.weight, + owner.eps, + owner.raw, + owner.output, + owner.use_pdl, + ) + + @staticmethod + def _bucket_flags(layer: M3DecoderLayer, num_tokens: int, hidden: int): + """(bucket workspace, its 8-CTA call counters) of the boundary all-reduce for ``num_tokens``, or None.""" + workspace = m3_mc_allreduce.get_workspace(layer.mapping, 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(): + dispatch.capture_missed(path, num_tokens, "the multicast all-reduce workspace") + else: + dispatch.missed(path, num_tokens, "no multicast all-reduce bucket holds the step") + + def _attn_push_target(self, layer: M3DecoderLayer, num_tokens: int, hidden: 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(layer, num_tokens, hidden) + if found is None: + self._no_bucket(dispatch.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, layer: M3DecoderLayer, num_tokens: int, hidden: 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 + (``M3MoE.prebuild_engines``), so it does.""" + from .ops.moe.m3_moe_engine import ar_push_desc + + found = self._bucket_flags(layer, num_tokens, hidden) + if found is None: + self._no_bucket(dispatch.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(): + dispatch.capture_missed( + dispatch.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, layer: M3DecoderLayer, use_mc: bool, x, residual, norm, x2=None, quant=False + ): + """All-reduce + (second addend) + residual + RMSNorm: the multicast kernel, else the built-in all-reduce.""" + if use_mc: + workspace = m3_mc_allreduce.get_workspace(layer.mapping, x.shape[-1]) + fused = ( + None + if workspace is None + else m3_mc_allreduce.m3_mc_allreduce_norm( + x, + residual.contiguous(), + norm.weight, + norm.variance_epsilon, + workspace, + use_pdl=get_env_enable_pdl(), + x2=x2, + quant=quant, + ) + ) + if fused is not None: + return fused + dispatch.missed( + dispatch.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 layer.native_allreduce_norm(x, residual, norm) + + def _prepushed(self, layer, x, residual, norm, quant, prepushed, x2=None): + workspace = m3_mc_allreduce.get_workspace(layer.mapping, x.shape[-1]) + fused = m3_mc_allreduce.m3_mc_allreduce_norm( + x, + residual.contiguous(), + norm.weight, + norm.variance_epsilon, + workspace, + use_pdl=get_env_enable_pdl(), + x2=x2, + quant=quant, + prepushed=prepushed, + ) + if fused is None: + raise RuntimeError( + f"layer {layer.layer_idx}: rows were pushed into an all-reduce that did not run" + ) + return fused + + def _layer( + self, + layer: M3DecoderLayer, + plan: DecodePlan, + position_ids, + hidden_states, + residual, + attn_metadata, + kwargs, + ): + num_tokens = hidden_states.shape[0] + hidden = hidden_states.shape[-1] + # 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: + dispatch.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 dispatch.ATTN_PUSH.covers(num_tokens) + armed = None + if push: + armed = self._attn_push_target(layer, num_tokens, hidden) + m3_prequant.AR_PUSH.arm(armed) + hidden_states = layer.self_attn( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + all_reduce_params=AllReduceParams(enable_allreduce=False), + **kwargs, + ) + if push: + m3_prequant.AR_PUSH.target = None + if push and m3_prequant.AR_PUSH.take_pushed(hidden_states): + hidden_states, residual = self._prepushed( + layer, hidden_states, residual, layer.post_attention_layernorm, False, AR_PUSH_MODE + ) + else: + if armed is not None: + dispatch.missed( + dispatch.ATTN_PUSH.path, + num_tokens, + "the o_proj did not push its rows into the armed boundary all-reduce", + ) + hidden_states, residual = self._boundary( + layer, plan.attn_mc_ar, hidden_states, residual, layer.post_attention_layernorm + ) + + moe = layer.block_sparse_moe + if moe is None: + out = layer.mlp( + hidden_states, final_all_reduce_params=AllReduceParams(enable_allreduce=False) + ) + return self._boundary(layer, plan.attn_mc_ar, out, residual, layer.next_layer_layernorm) + + norm = layer.next_layer_layernorm + quant = plan.qkv_prequant and layer.next_is_decoder_layer + path = dispatch.moe_path(num_tokens) + moe_push = None + if path == dispatch.MOE_ENGINE.path and moe.shared_experts is not None: + # The engine itself runs inside the MoE forward; looked up (or found missing) here first so a push is + # armed only for an engine that will take it. + engine = moe.engine(num_tokens, shared=True) + if engine is not None and plan.moe_mc_ar and decode_graph: + moe_push = self._moe_push_targets(layer, num_tokens, hidden) + if moe_push is not None: + m3_prequant.MOE_PUSH.arm(*moe_push) + elif path == dispatch.MOE16_ENGINE.path and decode_graph: + # Router, routed and shared experts in one kernel; its bf16 sum goes to the boundary all-reduce. + engine = moe.moe16_engine(num_tokens) + if engine is not None: + out = m3_moe16_engine(hidden_states.contiguous(), engine) + return self._boundary(layer, plan.moe_mc_ar, out, residual, norm, quant=quant) + if ( + path in (dispatch.MOE16_ENGINE.path, dispatch.MC_FINALIZE.path) + and plan.moe_mc_ar + and moe.shared_experts is not None + ): + if self._bucket_flags(layer, num_tokens, hidden) is None: + self._no_bucket(dispatch.MC_FINALIZE.path, num_tokens) + else: + shared, fc2, scales, perm = moe( + hidden_states, + attn_metadata, + final_all_reduce_params=AllReduceParams(enable_allreduce=False), + do_finalize=False, + ) + 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.weight, + norm.variance_epsilon, + m3_mc_allreduce.get_workspace(layer.mapping, 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" + ) + return fused + parts = moe( + hidden_states, + attn_metadata, + final_all_reduce_params=AllReduceParams(enable_allreduce=False), + return_parts=True, + ) + shared, routed = parts if isinstance(parts, tuple) else (parts, None) + if moe_push is not None: + m3_prequant.MOE_PUSH.disarm() + if routed is None: + # The engine folded the shared expert in and pushed routed + shared into set A. + folded_pushed = m3_prequant.MOE_PUSH.engine_pushed == shared.data_ptr() + m3_prequant.MOE_PUSH.engine_pushed = m3_prequant.MOE_PUSH.shared_pushed = None + if folded_pushed: + return self._prepushed(layer, shared, residual, norm, quant, 2) + else: + routed_pushed, shared_pushed = m3_prequant.MOE_PUSH.take_pushed(shared, routed) + if routed_pushed and shared_pushed: + return self._prepushed(layer, shared, residual, norm, quant, 3, x2=routed) + if routed_pushed or shared_pushed: + raise RuntimeError("only one MoE part pushed into the boundary all-reduce") + dispatch.missed( + dispatch.MOE_PUSH.path, + num_tokens, + "the MoE did not push its output into the armed boundary all-reduce", + ) + return self._boundary(layer, plan.moe_mc_ar, shared, residual, norm, x2=routed, quant=quant) + + +class ModelingV2MiniMaxM3Nvfp4Sm103Tp4(M3SpecDecForwardMixin, MiniMaxM3ForCausalLM): + """MiniMax-M3 NVFP4 / sm_103 / tp4, text only, with one-model Eagle3 or no speculation. + + A subclass of the built-in text model shell for its loader, its KV-cache and + transceiver preferences and its one-engine speculative plumbing; the text + model it wraps, the forward and the speculative worker are this target's. + """ + + 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 (dispatch.py). + dispatch.set_require(resolved_mode(type(self)) is ModelingV2Mode.REQUIRE) + if is_minimax_m3_vl_config(model_config.pretrained_config): + model_config = get_text_model_config(model_config) + 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 + # The one-engine shell's constructor, past the built-in one that builds the built-in text model. + SpecDecOneEngineForCausalLM.__init__( + self, MiniMaxM3Tp4ModelingV2(model_config), model_config + ) + self.runtime_hooks, self._spec_fold = make_runtime_hooks(self) + if self.spec_worker is not None: + self.spec_worker = M3Eagle3Worker.adopt(self.spec_worker) + worker = self.spec_worker + for layer in self.model.layers: + layer.capture = worker.capture_hidden_states + + def load_weights(self, weights: Dict, *args, **kwargs) -> None: + _weights.load(self, weights, *args, **kwargs) + + def setup_aliases(self) -> None: + """Chain each decoder layer's next_layer_layernorm for the fused boundaries.""" + super().setup_aliases() + layers = self.model.layers + for idx, layer in enumerate(layers): + layer.next_is_decoder_layer = idx < len(layers) - 1 + + def cache_derived_state(self) -> None: + super().cache_derived_state() + self.model.build_targets() + if self.spec_worker is not None: + # The fused verify's coded target head, from the loaded weights, before the executor's warmup. + start = time.perf_counter() + if self.spec_worker.prebuild_target_head(self.lm_head): + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] coded target head built before warmup: %.1f s" + % (time.perf_counter() - start) + ) + # The fused draft loop's head copy, from the loaded draft weights, before the executor's warmup. + start = time.perf_counter() + if self.spec_worker.prebuild_draft_head(self.draft_model): + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] fused draft head built before warmup: %.1f s" + % (time.perf_counter() - start) + ) + + def forward(self, attn_metadata: AttentionMetadata, *args, **kwargs): + """The fused-verify one-engine forward (``M3SpecDecForwardMixin``), text only.""" + multimodal_params: List[Any] = kwargs.pop("multimodal_params", None) or [] + if any( + param is not None and getattr(param, "multimodal_data", None) + for param in multimodal_params + ): + raise ValueError( + "this MiniMax-M3 target is text only; the request carries multimodal input" + ) + return super().forward(attn_metadata, *args, **kwargs) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/ops/attention/m3_attn_c1.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_attn_c1.py new file mode 100644 index 000000000000..625a7d5d4edc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_attn_c1.py @@ -0,0 +1,280 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The c1 kernel: fused MSA decode attention + o_proj GEMV in one 96-CTA launch, for a one-request, 6-token +verify layer.""" + +from __future__ import annotations + +from typing import NamedTuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_attn_mega import HIDDEN, check_handoff_at_rest, workspace_parity +from .m3_fused_decode_attn import ( + HEAD_DIM, + MXFP8_BLOCK, + NUM_HEADS, + check_attention_inputs, + mx_gemm, + mxfp8_rows, + reference_attention, + unswizzle_scales, +) + +#: 6 clusters of 16 CTAs, each an attention CTA and an o_proj unit: ``flags[32 + CTA]`` are the call counters. +GRID = 96 + + +class C1State(NamedTuple): + """The c1 kernel's results: the attention's bf16 output, the caches after the insert, and the o_proj output.""" + + output: torch.Tensor + kv_data: torch.Tensor + kv_scale: torch.Tensor + idx_k: torch.Tensor + y: torch.Tensor + + +class M3AttnC1(OpWrapper): + """The attention of a one-request verify step of 6 query tokens and its o_proj, in one 96-CTA launch. + + The fused attention of ``m3_fused_decode_attn`` on the packed QKV rows (bf16 ``output``, MXFP8 hand-off into the + call's parity of ``mxfp8_data`` / ``mxfp8_scale``), and the MXFP8 o_proj GEMV into ``y``: every attention CTA + is also an o_proj unit (48 row tiles x 2 K splits). ``flags[32 + CTA]`` are the per-CTA call counters; the + workspace is allocated once and never reset. With ``ar_mode`` 2, ``y`` is also pushed into the next multicast + all-reduce's exchange at ``ar_workspace`` / ``ar_flags``. + + Returns None; the cells read ``output``, ``y``, the hand-off and the caches. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="Eagle3 verify of 5 drafts (qlen 6) a few blocks in; both call parities (two calls, one workspace)", + spec=dict(seq_len=300, calls=2), + ), + Cell( + why="the full 16-block window: the six tokens fill the last rows of the last page", + spec=dict(seq_len=2048, calls=1), + ), + Cell( + why="the six tokens straddle a page boundary (positions 126..131)", + spec=dict(seq_len=132, calls=1), + ), + Cell(why="the first six tokens of a request", spec=dict(seq_len=6, calls=1)), + ) + + note = """ + Compiled for the one shape the target launches: one request of 6 query tokens (geometry (6, 1, 1), 96 CTAs), + the MXFP8 hand-off. The all-reduce push (`ar_mode` 2) needs a multicast workspace and is not driven here: the + cells run `ar_mode` 0, where `y` is the o_proj output and nothing else is written. + """ + + def __call__( + self, + packed: torch.Tensor, + flags: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + o_weight: torch.Tensor, + o_weight_scale: torch.Tensor, + y: torch.Tensor, + ar_workspace: int, + ar_flags: int, + ar_rank: int, + ar_world: int, + ar_max_tokens: int, + ar_mode: int, + ) -> None: + torch.ops.trtllm.minimax_m3_attn_c1( + packed, + flags, + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + eps, + sm_scale, + decode_query_len, + tokens_per_cluster, + blocks_per_cta, + output, + mxfp8_data, + mxfp8_scale, + use_pdl, + o_weight, + o_weight_scale, + y, + ar_workspace, + ar_flags, + ar_rank, + ar_world, + ar_max_tokens, + ar_mode, + ) + + def reference( + self, + packed: torch.Tensor, + flags: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + o_weight: torch.Tensor, + o_weight_scale: torch.Tensor, + y: torch.Tensor, + ar_workspace: int, + ar_flags: int, + ar_rank: int, + ar_world: int, + ar_max_tokens: int, + ar_mode: int, + ) -> C1State: + """The composition: the fused attention's reference on ``packed``, then + ``y = bf16(mxfp8(output) @ dequant(W_o)^T)`` (fp64 accumulation, one rounding). + + The workspace, the output buffers, the geometry, ``use_pdl`` and the all-reduce arguments are ignored. + """ + attn = reference_attention( + packed, + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + eps, + sm_scale, + decode_query_len, + ) + q8, ex = mxfp8_rows(attn.output) + width = NUM_HEADS * HEAD_DIM + out_y = mx_gemm( + q8, + ex, + o_weight.view(torch.uint8), + unswizzle_scales(o_weight_scale, HIDDEN, width // MXFP8_BLOCK), + ) + return C1State(attn.output, attn.kv_data, attn.kv_scale, attn.idx_k, out_y) + + def is_valid( + self, + packed: torch.Tensor, + flags: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + o_weight: torch.Tensor, + o_weight_scale: torch.Tensor, + y: torch.Tensor, + ar_workspace: int, + ar_flags: int, + ar_rank: int, + ar_world: int, + ar_max_tokens: int, + ar_mode: int, + ) -> None: + check_attention_inputs( + packed.shape[0], + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + decode_query_len, + ) + check_handoff_at_rest(mxfp8_data, mxfp8_scale, workspace_parity(flags, GRID)) + assert ar_mode != 2 or (ar_workspace != 0 and ar_flags != 0), ( + "ar_mode 2 needs the exchange's addresses" + ) + + +m3_attn_c1 = M3AttnC1() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_attn_mega.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_attn_mega.py new file mode 100644 index 000000000000..3445c8c65cbd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_attn_mega.py @@ -0,0 +1,340 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The attention megakernel: QKV GEMV + fused MSA decode attention + o_proj GEMV in one launch, for a one-request, +one-token decode layer.""" + +from __future__ import annotations + +from typing import NamedTuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_fused_decode_attn import ( + HEAD_DIM, + MXFP8_BLOCK, + NUM_HEADS, + PACKED_WIDTH, + check_attention_inputs, + mx_gemm, + mxfp8_rows, + reference_attention, + unswizzle_scales, +) + +#: The layer input (hidden size) and the o_proj output width at tp4. +HIDDEN = 6144 +#: Persistent grid of the megakernel (7 clusters of 16): ``flags[32 + CTA]`` are its per-CTA call counters. +GRID = 112 +#: Rows of each call-parity buffer. +MAX_TOKENS = 8 + + +class MegaState(NamedTuple): + """The megakernel's results: the packed QKV rows of the call, the attention's bf16 output, the caches after the + insert, and the o_proj output.""" + + packed: torch.Tensor + output: torch.Tensor + kv_data: torch.Tensor + kv_scale: torch.Tensor + idx_k: torch.Tensor + y: torch.Tensor + + +def workspace_parity(flags: torch.Tensor, grid: int) -> int: + """The call parity every CTA of the grid will use (their counters move together), checked.""" + counters = flags[32 : 32 + grid] + assert bool((counters == counters[0]).all()), ( + "the per-CTA call counters disagree: the workspace is corrupt" + ) + return int(counters[0]) & 1 + + +def check_handoff_at_rest(mxfp8_data: torch.Tensor, mxfp8_scale: torch.Tensor, parity: int) -> None: + """The MXFP8 hand-off's parity buffer this call writes must hold the 0xFF "not yet" marker (Lamport).""" + assert bool((mxfp8_data.view(torch.uint8).reshape(2, -1)[parity] == 0xFF).all()) and bool( + (mxfp8_scale.view(torch.uint8).reshape(2, -1)[parity] == 0xFF).all() + ), ( + "the MXFP8 hand-off's buffer for this call is not at rest (0xFF): the o_proj units would read stale codes" + ) + + +class M3AttnMega(OpWrapper): + """The whole attention half of a decode layer for one request and one query token, in one launch. + + The MXFP8 QKV GEMV of the layer input (as the preceding all-reduce hands it over: E4M3 rows + 128x4-swizzled + UE8M0 scales) into the call's parity buffer of ``qkv_out``, the fused attention of ``m3_fused_decode_attn`` on + those rows (bf16 ``output``, MXFP8 hand-off into the call's parity of ``mxfp8_data`` / ``mxfp8_scale``), and the + MXFP8 o_proj GEMV into ``y``. ``flags`` holds the per-tile ready counters and the per-CTA call counters; the + workspace is allocated once and never reset. With ``ar_mode`` 2, ``y`` is also pushed into the next multicast + all-reduce's exchange at ``ar_workspace`` / ``ar_flags``. + + Returns None; the cells read ``qkv_out``, ``output``, ``y``, the hand-off and the caches. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode of one request a few blocks in; both call parities (two calls on one workspace)", + spec=dict(seq_len=300, calls=2), + ), + Cell( + why="the full 16-block window, the new token on the last row of the last page", + spec=dict(seq_len=2048, calls=1), + ), + Cell( + why="the new token opens a page (position 1024): its own page's other rows are past the length", + spec=dict(seq_len=1025, calls=1), + ), + Cell(why="the first token of a request: one attended row", spec=dict(seq_len=1, calls=1)), + ) + + note = """ + Compiled for the one shape the target launches: one request, one query token (geometry (1, 1, 1)), the MXFP8 + hand-off. The all-reduce push (`ar_mode` 2) needs a multicast workspace and is not driven here: the cells run + `ar_mode` 0, where `y` is the o_proj output and nothing else is written. + + The QKV rows are handed from the GEMV CTAs to the attention CTAs by polling for the 0xFFFF marker the + previous call left: a GEMV output whose bf16 bits are 0xFFFF (a NaN) would never be seen as written. Finite + inputs cannot produce it. + """ + + def __call__( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + qkv_weight: torch.Tensor, + qkv_weight_scale: torch.Tensor, + qkv_out: torch.Tensor, + flags: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + o_weight: torch.Tensor, + o_weight_scale: torch.Tensor, + y: torch.Tensor, + ar_workspace: int, + ar_flags: int, + ar_rank: int, + ar_world: int, + ar_max_tokens: int, + ar_mode: int, + ) -> None: + torch.ops.trtllm.minimax_m3_attn_mega( + xq, + xsf, + qkv_weight, + qkv_weight_scale, + qkv_out, + flags, + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + eps, + sm_scale, + decode_query_len, + tokens_per_cluster, + blocks_per_cta, + output, + mxfp8_data, + mxfp8_scale, + use_pdl, + o_weight, + o_weight_scale, + y, + ar_workspace, + ar_flags, + ar_rank, + ar_world, + ar_max_tokens, + ar_mode, + ) + + def reference( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + qkv_weight: torch.Tensor, + qkv_weight_scale: torch.Tensor, + qkv_out: torch.Tensor, + flags: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + o_weight: torch.Tensor, + o_weight_scale: torch.Tensor, + y: torch.Tensor, + ar_workspace: int, + ar_flags: int, + ar_rank: int, + ar_world: int, + ar_max_tokens: int, + ar_mode: int, + ) -> MegaState: + """The composition: ``packed = bf16(dequant(xq) @ dequant(W_qkv)^T)``, the fused attention's reference on + those rows, then ``y = bf16(mxfp8(output) @ dequant(W_o)^T)`` (fp64 accumulation, one rounding per GEMV). + + The workspace, the output buffers, the geometry, ``use_pdl`` and the all-reduce arguments are ignored. + """ + t = xq.shape[0] + xq_b, w_b = xq.view(torch.uint8), qkv_weight.view(torch.uint8) + packed = mx_gemm( + xq_b, + unswizzle_scales(xsf, t, HIDDEN // MXFP8_BLOCK), + w_b, + unswizzle_scales(qkv_weight_scale, PACKED_WIDTH, HIDDEN // MXFP8_BLOCK), + ) + attn = reference_attention( + packed, + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + eps, + sm_scale, + decode_query_len, + ) + q8, ex = mxfp8_rows(attn.output) + width = NUM_HEADS * HEAD_DIM + out_y = mx_gemm( + q8, + ex, + o_weight.view(torch.uint8), + unswizzle_scales(o_weight_scale, HIDDEN, width // MXFP8_BLOCK), + ) + return MegaState(packed, attn.output, attn.kv_data, attn.kv_scale, attn.idx_k, out_y) + + def is_valid( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + qkv_weight: torch.Tensor, + qkv_weight_scale: torch.Tensor, + qkv_out: torch.Tensor, + flags: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + o_weight: torch.Tensor, + o_weight_scale: torch.Tensor, + y: torch.Tensor, + ar_workspace: int, + ar_flags: int, + ar_rank: int, + ar_world: int, + ar_max_tokens: int, + ar_mode: int, + ) -> None: + check_attention_inputs( + xq.shape[0], + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + decode_query_len, + ) + assert xsf.numel() >= 128 * (HIDDEN // MXFP8_BLOCK), ( + "xsf must hold one 128-row tile of swizzled scales" + ) + parity = workspace_parity(flags, GRID) + assert bool((qkv_out.view(torch.int16)[parity] == -1).all()), ( + "qkv_out's buffer for this call is not at rest (0xFFFF): the attention would read stale QKV rows" + ) + check_handoff_at_rest(mxfp8_data, mxfp8_scale, parity) + assert ar_mode != 2 or (ar_workspace != 0 and ar_flags != 0), ( + "ar_mode 2 needs the exchange's addresses" + ) + + +m3_attn_mega = M3AttnMega() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_dense_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_dense_decode.py new file mode 100644 index 000000000000..38d5f1342c78 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_dense_decode.py @@ -0,0 +1,187 @@ +# 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 (``dispatch.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.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 (dispatch.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_103}) + + 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 = q.to(torch.float8_e4m3fn) + 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_103__tp4/ops/attention/m3_fused_decode_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_fused_decode_attn.py new file mode 100644 index 000000000000..9a882ea6b933 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_fused_decode_attn.py @@ -0,0 +1,643 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fused MSA short-sequence decode attention of this target's sparse layers, and the plain-torch +composition it fuses (shared with the megakernel and c1 entries, which wrap the same attention).""" + +from __future__ import annotations + +import math +from typing import NamedTuple, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Per-rank attention geometry at tp4: 16 query heads, one KV head, one index head, head_dim 128. +NUM_HEADS = 16 +HEAD_DIM = 128 +#: The packed QKV(+index) projection row: [Q x16 | K | V | idxQ | idxK] x 128. +PACKED_WIDTH = (NUM_HEADS + 4) * HEAD_DIM +#: Partial NeoX RoPE over the first 64 dims of every head. +ROTARY_DIM = 64 +#: One MSA block is one 128-token KV page; identity selection covers at most 16 of them. +PAGE = 128 +MAX_BLOCKS = 16 +MAX_SEQ_LEN = MAX_BLOCKS * PAGE +#: NVFP4 block: 16 elements per E4M3 scale. +NVFP4_BLOCK = 16 +#: MXFP8 block of the o_proj hand-off: 32 elements per UE8M0 scale. +MXFP8_BLOCK = 32 + +_E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +_E4M3_MAX = 448.0 + + +def e2m1_values(device) -> torch.Tensor: + """The 16 E2M1 codes' values (sign in bit 3).""" + pos = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=device) + return torch.cat([pos, -pos]) + + +def e4m3_bytes(x: torch.Tensor) -> torch.Tensor: + """Round fp32 to E4M3 the way ``__nv_fp8_e4m3(float)`` does (RN-even, saturating to +-448), as raw bytes.""" + return x.clamp(-_E4M3_MAX, _E4M3_MAX).to(torch.float8_e4m3fn).view(torch.uint8) + + +def e4m3_value(b: torch.Tensor) -> torch.Tensor: + return b.view(torch.float8_e4m3fn).float() + + +def e2m1_codes(x: torch.Tensor) -> torch.Tensor: + """``cvt.rn.satfinite.e2m1x2.f32`` per element: round-to-nearest-even onto {0, .5, 1, 1.5, 2, 3, 4, 6}, + saturating at 6, sign in bit 3 (a negative value that rounds to zero keeps its sign).""" + a = x.abs() + # Ties go to the even code: 0.25 -> 0, 0.75 -> 1.0, 1.25 -> 1.0, 1.75 -> 2.0, 2.5 -> 2.0, 3.5 -> 4.0, 5 -> 4. + code = ( + (a > 0.25).int() + + (a >= 0.75).int() + + (a > 1.25).int() + + (a >= 1.75).int() + + (a > 2.5).int() + + (a >= 3.5).int() + + (a > 5.0).int() + ) + return (code | (torch.signbit(x).int() << 3)).to(torch.uint8) + + +def nvfp4_quantize( + x: torch.Tensor, second_level_scale: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + """The MiniMax-M3 producer's NVFP4 recipe for rows of 128: (packed data [N, 64] uint8, E4M3 scales [N, 8]). + + Elements are rounded to bf16 first; each 16-element block's scale is ``e4m3(sls * amax / 6)`` and the codes are + ``e2m1(x / (scale / sls))``. Element ``2i`` is the low nibble of byte ``i``. + """ + n = x.shape[0] + xb = x.to(torch.bfloat16).float().view(n, HEAD_DIM // NVFP4_BLOCK, NVFP4_BLOCK) + sls = second_level_scale.float().reshape(()) + amax = xb.abs().amax(-1) + scale = e4m3_bytes(sls * (amax * (1.0 / 6.0))) + out_scale = torch.where( + amax != 0, 1.0 / (e4m3_value(scale) * (1.0 / sls)), torch.zeros_like(amax) + ) + codes = e2m1_codes(xb * out_scale[..., None]).view(n, HEAD_DIM) + data = codes[:, 0::2] | (codes[:, 1::2] << 4) + return data, scale + + +def nvfp4_dequantize(data: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """[N, 64] packed E2M1 + [N, 8] E4M3 block scales -> fp32 [N, 128] (the per-tensor scale not applied).""" + table = e2m1_values(data.device) + lo = table[(data & 0xF).long()] + hi = table[(data >> 4).long()] + values = torch.stack([lo, hi], dim=-1).reshape(data.shape[0], HEAD_DIM) + return values * e4m3_value(scale).repeat_interleave(NVFP4_BLOCK, dim=-1) + + +def v_scale_rows(scale_plane: torch.Tensor) -> torch.Tensor: + """A page's V scale plane in the 4x4 token-quad order [token // 4][column][token % 4] -> [128, 8] row-major.""" + return ( + scale_plane.reshape(PAGE // 4, HEAD_DIM // NVFP4_BLOCK, 4) + .permute(0, 2, 1) + .reshape(PAGE, -1) + ) + + +def v_scale_plane(rows: torch.Tensor) -> torch.Tensor: + """Inverse of :func:`v_scale_rows`: [128, 8] row-major scales -> the plane's [128, 8] storage.""" + return rows.reshape(PAGE // 4, 4, HEAD_DIM // NVFP4_BLOCK).permute(0, 2, 1).reshape(PAGE, -1) + + +def norm_rope( + x: torch.Tensor, weight: torch.Tensor, eps: float, cos: torch.Tensor, sin: torch.Tensor +) -> torch.Tensor: + """Per-head Gemma RMSNorm (``x * rsqrt(mean(x^2) + eps) * (1 + w)``) then NeoX RoPE over the first 64 dims. + + x: fp32 [N, 128]; cos / sin: fp32 [N, 32] for each row's position.""" + y = x * (torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + weight.float())) + x1, x2 = y[:, : ROTARY_DIM // 2], y[:, ROTARY_DIM // 2 : ROTARY_DIM] + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, y[:, ROTARY_DIM:]], dim=-1) + + +def swizzled_scale_offsets(rows: int, k_blocks: int, device) -> torch.Tensor: + """Byte offset of scale (row, k-block) in the 128x4-swizzled UE8M0 layout: [rows, k_blocks] int64.""" + r = torch.arange(rows, device=device)[:, None] + k = torch.arange(k_blocks, device=device)[None, :] + k_tiles = (k_blocks + 3) // 4 + return ( + (r // 128) * k_tiles * 512 + (k // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + k % 4 + ) + + +def mxfp8_rows(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """The merge's MXFP8 hand-off recipe on bf16 rows [T, K]: (E4M3 bytes [T, K], UE8M0 exponents [T, K / 32]). + + Per 32-element block: ``e = ceil_exp(amax / 448)`` (biased, at most 254), values ``e4m3(x * 2^(127 - e))``.""" + t, k = x.shape + xf = x.float().view(t, k // MXFP8_BLOCK, MXFP8_BLOCK) + amax = xf.abs().amax(-1) + bits = (amax * (1.0 / _E4M3_MAX)).view(torch.int32) + ex = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).int() + ex = ex.clamp(max=254) + scale = torch.where(amax != 0, ((254 - ex) << 23).view(torch.float32), torch.zeros_like(amax)) + q8 = e4m3_bytes(xf * scale[..., None]).view(t, k) + return q8, ex.to(torch.uint8) + + +def mxfp8_dequantize(q8: torch.Tensor, ex: torch.Tensor) -> torch.Tensor: + """E4M3 bytes [T, K] with UE8M0 exponents [T, K / 32] -> fp64 values.""" + scale = torch.exp2(ex.double() - 127.0).repeat_interleave(MXFP8_BLOCK, dim=-1) + return e4m3_value(q8).double() * scale + + +def unswizzle_scales(swizzled: torch.Tensor, rows: int, k_blocks: int) -> torch.Tensor: + """The [rows, k_blocks] UE8M0 bytes of a 128x4-swizzled scale buffer.""" + flat = swizzled.reshape(-1).view(torch.uint8) + return flat[swizzled_scale_offsets(rows, k_blocks, flat.device)] + + +def mx_gemm( + a: torch.Tensor, a_exp: torch.Tensor, w: torch.Tensor, w_exp: torch.Tensor +) -> torch.Tensor: + """bf16 ``dequant(a) @ dequant(w)^T`` of MXFP8 operands (E4M3 bytes, row-major UE8M0 exponents), accumulated in + fp64 and rounded once.""" + return (mxfp8_dequantize(a, a_exp) @ mxfp8_dequantize(w, w_exp).T).to(torch.bfloat16) + + +class AttnState(NamedTuple): + """What the fused attention leaves behind: its bf16 output and the three caches after the new tokens' insert.""" + + output: torch.Tensor + kv_data: torch.Tensor + kv_scale: torch.Tensor + idx_k: torch.Tensor + + +def query_rows(packed, position_ids, q_norm_weight, eps, rotary_cos_sin) -> torch.Tensor: + """fp32 [T, 16, 128]: the query heads of packed rows after Gemma RMSNorm + NeoX RoPE (before E4M3).""" + rows = packed.float().view(packed.shape[0], NUM_HEADS + 4, HEAD_DIM) + pos_ids = position_ids.long() + cos = rotary_cos_sin[pos_ids, 0].float() + sin = rotary_cos_sin[pos_ids, 1].float() + return torch.stack( + [norm_rope(rows[:, h], q_norm_weight, eps, cos, sin) for h in range(NUM_HEADS)], dim=1 + ) + + +def attend_cached( + q, + kv_data, + kv_scale, + block_table, + seq_lens, + k_global_scale, + v_global_scale, + sm_scale, + decode_query_len, +) -> torch.Tensor: + """bf16 [T, 2048]: the identity-selection decode of fp32 queries ``q`` [T, 16, 128] over caches that already + hold the step's tokens. q is rounded to E4M3 (no scale), the operand the kernel's tensor core reads; request + ``r``'s token ``t`` sits at ``pos = seq_lens[r] - qlen + t`` and attends positions ``[0, pos]`` through + ``block_table``: scores ``q . dequant(K) * sm_scale * k_global_scale``, softmax, values + ``dequant(V) * v_global_scale``. fp64 throughout, one bf16 rounding at the end (the kernel's split-K partials + and their LSE merge are exact algebra).""" + q8 = e4m3_value(e4m3_bytes(q)).double() + t_total = q.shape[0] + qlen = int(decode_query_len) + out = torch.empty(t_total, NUM_HEADS, HEAD_DIM, dtype=torch.float64, device=q.device) + k_gs = float(k_global_scale.reshape(-1)[0]) + v_gs = float(v_global_scale.reshape(-1)[0]) + for r, seq_len in enumerate(seq_lens.tolist()): + pages = block_table[r, : math.ceil(seq_len / PAGE)].tolist() + kd = torch.cat([nvfp4_dequantize(kv_data[p, 0, 0], kv_scale[p, 0, 0]) for p in pages]) + vd = torch.cat( + [nvfp4_dequantize(kv_data[p, 1, 0], v_scale_rows(kv_scale[p, 1, 0])) for p in pages] + ) + kd, vd = kd.double() * k_gs, vd.double() * v_gs + for t in range(qlen): + row = r * qlen + t + n = seq_len - qlen + t + 1 + p = torch.softmax(q8[row] @ kd[:n].T * sm_scale, dim=-1) + out[row] = p @ vd[:n] + return out.to(torch.bfloat16).view(t_total, NUM_HEADS * HEAD_DIM) + + +def reference_attention( + packed, + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + eps, + sm_scale, + decode_query_len, +) -> AttnState: + """The composition the fused kernel replaces, in plain torch, on copies of the caches. + + 1. Per token, q heads / k / index-k: Gemma RMSNorm + NeoX RoPE at ``position_ids``; v as is. + 2. k / v -> NVFP4 (``kv_quant_scale[1]`` / ``[2]``) into the token's ``out_cache_loc`` slot (K scales linear, V + scales in 4x4 token quads); index-k -> E4M3 into the index cache. + 3. :func:`attend_cached` over the caches after the insert. + """ + kv_data, kv_scale, idx_k = kv_data.clone(), kv_scale.clone(), idx_k.clone() + idx_bytes = idx_k.view(torch.uint8) + t_total = packed.shape[0] + rows = packed.float().view(t_total, NUM_HEADS + 4, HEAD_DIM) + pos_ids = position_ids.long() + cos = rotary_cos_sin[pos_ids, 0].float() + sin = rotary_cos_sin[pos_ids, 1].float() + q = query_rows(packed, position_ids, q_norm_weight, eps, rotary_cos_sin) + k = norm_rope(rows[:, NUM_HEADS], k_norm_weight, eps, cos, sin) + v = rows[:, NUM_HEADS + 1] + ik = norm_rope(rows[:, NUM_HEADS + 3], index_k_norm_weight, eps, cos, sin) + k_data, k_sf = nvfp4_quantize(k, kv_quant_scale[1]) + v_data, v_sf = nvfp4_quantize(v, kv_quant_scale[2]) + ik_bytes = e4m3_bytes(ik) + for t, slot in enumerate(out_cache_loc[:t_total].tolist()): + page, row = divmod(slot, PAGE) + kv_data[page, 0, 0, row] = k_data[t] + kv_data[page, 1, 0, row] = v_data[t] + kv_scale[page, 0, 0, row] = k_sf[t] + vrows = v_scale_rows(kv_scale[page, 1, 0]) + vrows[row] = v_sf[t] + kv_scale[page, 1, 0] = v_scale_plane(vrows) + idx_bytes[page, 0, row] = ik_bytes[t] + output = attend_cached( + q, + kv_data, + kv_scale, + block_table, + seq_lens, + k_global_scale, + v_global_scale, + sm_scale, + decode_query_len, + ) + return AttnState(output, kv_data, kv_scale, idx_k) + + +def check_attention_inputs( + packed_rows, + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + decode_query_len, +) -> None: + """The inputs the attention kernels accept and answer wrongly (shared by the three entries' ``is_valid``).""" + assert ( + kv_data.dim() == 5 + and tuple(kv_data.shape[1:]) == (2, 1, PAGE, HEAD_DIM // 2) + and kv_data.stride(4) == 1 + and kv_data.stride(3) == HEAD_DIM // 2 + ), "kv_data must be [pages, 2, 1, 128, 64] with contiguous head blocks" + assert ( + kv_scale.dim() == 5 + and tuple(kv_scale.shape[1:]) == (2, 1, PAGE, HEAD_DIM // NVFP4_BLOCK) + and kv_scale.stride(4) == 1 + and kv_scale.stride(3) == HEAD_DIM // NVFP4_BLOCK + ), "kv_scale must be [pages, 2, 1, 128, 8] with contiguous head blocks" + assert ( + idx_k.dim() == 4 and tuple(idx_k.shape[1:]) == (1, PAGE, HEAD_DIM) and idx_k.stride(3) == 1 + ), "idx_k must be [pages, 1, 128, 128] with contiguous rows" + for name, w in ( + ("q_norm_weight", q_norm_weight), + ("k_norm_weight", k_norm_weight), + ("index_k_norm_weight", index_k_norm_weight), + ): + assert w.is_contiguous(), f"{name} must be contiguous (read as 4-element vectors)" + assert rotary_cos_sin.is_contiguous() and tuple(rotary_cos_sin.shape[1:]) == ( + 2, + ROTARY_DIM // 2, + ), "rotary_cos_sin must be contiguous [max_pos, 2, 32]" + for name, t in ( + ("out_cache_loc", out_cache_loc), + ("position_ids", position_ids), + ("seq_lens", seq_lens), + ): + assert t.is_contiguous(), f"{name} must be contiguous" + qlen = int(decode_query_len) + num_reqs = seq_lens.shape[0] + assert packed_rows == num_reqs * qlen, "one row per (request, query token)" + assert position_ids.numel() == packed_rows, "one position per token" + lens = seq_lens.tolist() + assert all(qlen <= n <= MAX_SEQ_LEN for n in lens), ( + f"every seq_len must be in [decode_query_len, {MAX_SEQ_LEN}]: identity selection covers 16 blocks of 128, " + "a longer request is silently cut to them" + ) + assert block_table.shape[1] >= math.ceil(max(lens) / PAGE), ( + "block_table must hold every attended block (missing columns read page 0)" + ) + assert int(position_ids.max()) < rotary_cos_sin.shape[0], "position past the rotary table" + table = block_table.tolist() + slots = out_cache_loc.tolist() + for r, n in enumerate(lens): + for t in range(qlen): + pos = n - qlen + t + want = table[r][pos // PAGE] * PAGE + pos % PAGE + assert slots[r * qlen + t] == want, ( + "out_cache_loc must be the slot block_table gives the token's position seq_len - qlen + t: the " + "attention reads that page's other rows from the cache it loaded before the insert" + ) + + +class M3FusedDecodeAttn(OpWrapper): + """One launch after the packed QKV GEMM of a sparse layer in short-sequence decode. + + Fuses q / k / index-k Gemma RMSNorm + NeoX partial RoPE, the NVFP4 K/V and FP8 index-K cache insert, the + block-sparse decode over every attended block (selection is the identity: at most 16 blocks of 128 tokens) and + the split-K LSE merge, writing the bf16 output and its MXFP8 copy for the o_proj GEMM. + + ``(tokens_per_cluster, blocks_per_cta)`` is the launch geometry; the target passes + ``kernels.m3_fused_attn.default_geometry(qlen, num_requests)``. Returns None: the cells read ``output``, + ``mxfp8_data`` / ``mxfp8_scale`` and the three caches. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # Every compiled geometry, each at a step shape the target routes to it (decode_query_len x requests, at most + # 16 tokens). seq_lens is a per-request list; positions are seq_len - qlen + t. + CELLS: tuple[Cell, ...] = ( + Cell( + why="plain decode, one request, a few blocks: geometry (1, 1, 1)", + spec=dict(qlen=1, seq_lens=(300,), geometry=(1, 1)), + ), + Cell( + why="plain decode, 6 requests: the widest one-block-per-CTA batch; full 16 blocks and a page's 1st row", + spec=dict(qlen=1, seq_lens=(2048, 129, 1, 700, 1500, 256), geometry=(1, 1)), + ), + Cell( + why="plain decode, 8 requests: KV blocks paired per CTA (8-CTA clusters), geometry (1, 1, 2)", + spec=dict(qlen=1, seq_lens=(2048, 33, 128, 129, 1024, 1900, 640, 7), geometry=(1, 2)), + ), + Cell( + why="plain decode at the 16-token cap: 16 paired-block clusters", + spec=dict(qlen=1, seq_lens=tuple(range(97, 2048, 125))[:16], geometry=(1, 2)), + ), + Cell( + why="Eagle3 verify of 1 draft (qlen 2), 2 requests: geometry (2, 1, 1); a draft crossing a page", + spec=dict(qlen=2, seq_lens=(129, 900), geometry=(1, 1)), + ), + Cell( + why="qlen 2 at 8 requests: geometry (2, 1, 2)", + spec=dict(qlen=2, seq_lens=(2, 2048, 130, 257, 640, 1000, 1500, 64), geometry=(1, 2)), + ), + Cell( + why="qlen 3 at 5 requests (15 tokens): geometry (3, 1, 1)", + spec=dict(qlen=3, seq_lens=(3, 2048, 130, 385, 777), geometry=(1, 1)), + ), + Cell( + why="qlen 4 at 4 requests: geometry (4, 1, 1)", + spec=dict(qlen=4, seq_lens=(4, 2048, 131, 1030), geometry=(1, 1)), + ), + Cell( + why="qlen 5 at 3 requests: geometry (5, 1, 1)", + spec=dict(qlen=5, seq_lens=(5, 1100, 2048), geometry=(1, 1)), + ), + Cell( + why="qlen 6, one request (the c1 shape when c1 does not apply): geometry (6, 1, 1)", + spec=dict(qlen=6, seq_lens=(1283,), geometry=(1, 1)), + ), + Cell( + why="qlen 6 at 2 requests: geometry (6, 1, 2)", + spec=dict(qlen=6, seq_lens=(6, 2048), geometry=(1, 2)), + ), + Cell( + why="qlen 7, two requests: geometry (7, 1, 1)", + spec=dict(qlen=7, seq_lens=(132, 2000), geometry=(1, 1)), + ), + Cell( + why="qlen 8 at 2 requests (the 16-token cap): geometry (8, 1, 1)", + spec=dict(qlen=8, seq_lens=(8, 1029), geometry=(1, 1)), + ), + # Every other (query rows per request, requests) the decode dispatch sends here (dispatch.FUSED_DECODE_ATTN): + # lengths spread over the 16-block window, the geometry the target picks for the shape. + *( + Cell( + why=f"qlen {qlen} at {reqs} requests: a decode shape between the ones above", + spec=dict( + qlen=qlen, + seq_lens=tuple( + qlen + (r * 389 + 97 * qlen) % (2049 - qlen) for r in range(reqs) + ), + geometry=geometry, + ), + ) + for qlen, reqs, geometry in ( + (1, 2, (1, 1)), + (1, 3, (1, 1)), + (1, 4, (1, 1)), + (1, 5, (1, 1)), + (1, 7, (1, 2)), + (1, 9, (1, 2)), + (1, 10, (1, 2)), + (1, 11, (1, 2)), + (1, 12, (1, 2)), + (1, 13, (1, 2)), + (1, 14, (1, 2)), + (1, 15, (1, 2)), + (2, 1, (1, 1)), + (2, 3, (1, 1)), + (2, 4, (1, 2)), + (2, 5, (1, 2)), + (2, 6, (1, 2)), + (2, 7, (1, 2)), + (3, 1, (1, 1)), + (3, 2, (1, 1)), + (3, 3, (1, 1)), + (3, 4, (1, 1)), + (4, 1, (1, 1)), + (4, 2, (1, 1)), + (4, 3, (1, 1)), + (5, 1, (1, 1)), + (5, 2, (1, 1)), + (7, 1, (1, 1)), + (8, 1, (1, 1)), + ) + ), + ) + + note = """ + Only the variants the target launches are compiled: the geometries above, each with the MXFP8 hand-off (the + target always passes it). Q is quantized to E4M3 with no scale (`kv_quant_scale[0]` is not read): its operand + on the tensor core is the E4M3 value, which `reference` reproduces. + + Not certified: a KV pool whose pages interleave other layers (any `stride(0)` / `stride(1)` that is a multiple + of 16 is accepted; the cells use a contiguous pool), `position_ids` other than `seq_len - qlen + t` (RoPE reads + `position_ids`, placement reads `seq_lens`; the target passes equal values), and `use_pdl=False` timing. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + ) -> None: + torch.ops.trtllm.minimax_m3_fused_decode_attn( + packed, + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + eps, + sm_scale, + decode_query_len, + tokens_per_cluster, + blocks_per_cta, + output, + mxfp8_data, + mxfp8_scale, + use_pdl, + ) + + def reference( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + ) -> AttnState: + """The bf16 output and the caches after the insert, as :func:`reference_attention` computes them. + + The geometry, the output buffers and ``use_pdl`` do not change what is computed and are ignored; the + MXFP8 hand-off is :func:`mxfp8_rows` of the bf16 output, which the test checks on the op's own output. + """ + return reference_attention( + packed, + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + kv_quant_scale, + k_global_scale, + v_global_scale, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + eps, + sm_scale, + decode_query_len, + ) + + def is_valid( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + idx_k: torch.Tensor, + out_cache_loc: torch.Tensor, + position_ids: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_quant_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + eps: float, + sm_scale: float, + decode_query_len: int, + tokens_per_cluster: int, + blocks_per_cta: int, + output: torch.Tensor, + mxfp8_data: torch.Tensor, + mxfp8_scale: torch.Tensor, + use_pdl: bool, + ) -> None: + check_attention_inputs( + packed.shape[0], + kv_data, + kv_scale, + idx_k, + out_cache_loc, + position_ids, + block_table, + seq_lens, + q_norm_weight, + k_norm_weight, + index_k_norm_weight, + rotary_cos_sin, + decode_query_len, + ) + k_blocks = output.shape[1] // MXFP8_BLOCK + needed = math.ceil(packed.shape[0] / 128) * ((k_blocks + 3) // 4) * 512 + assert mxfp8_scale.numel() >= needed, ( + "mxfp8_scale must hold the swizzled scales of every output row tile" + ) + + +m3_fused_decode_attn = M3FusedDecodeAttn() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_index_score.py new file mode 100644 index 000000000000..e7e57c10bde2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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 ``kernels/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 ...kernels.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_103}) + + 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_103__tp4/ops/attention/m3_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_producer.py new file mode 100644 index 000000000000..957a1b82eef0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_producer.py @@ -0,0 +1,197 @@ +# 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_103}) + + 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, + ) -> tuple[torch.Tensor, torch.Tensor]: + 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, + ) -> 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, + ) -> 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_103__tp4/ops/attention/m3_select_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_select_long.py new file mode 100644 index 000000000000..91c1817f1610 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_select_long.py @@ -0,0 +1,172 @@ +# 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_103}) + + 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, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_select_long( + scores, n_valid_blocks, out, init_blocks, local_blocks, use_pdl, early_trigger + ) + 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, + ) -> 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`` and + ``early_trigger`` 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, + ) -> 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_103__tp4/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode.py new file mode 100644 index 000000000000..3c59454dac0a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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 ...kernels.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_103}) + + 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_103__tp4/ops/attention/m3_sparse_decode_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode_lc.py new file mode 100644 index 000000000000..2dfeba5cbd0a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode_lc.py @@ -0,0 +1,395 @@ +# 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_103}) + + 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, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_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, + early_trigger, + ) + 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, + ) -> torch.Tensor: + """The BF16 ``out`` [T, 2048] the op is supposed to write: ``sparse_attention_reference`` in float64, + rounded once. ``linear_softmax``, ``use_pdl`` and ``early_trigger`` 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, + ) -> 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_103__tp4/ops/attention/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/msa_kv_update.py new file mode 100644 index 000000000000..ad05b934d4a9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp4/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_mc_allreduce_norm.py new file mode 100644 index 000000000000..a2b74d268d09 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_mc_allreduce_norm.py @@ -0,0 +1,277 @@ +# 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, +) -> tuple[torch.Tensor, ...]: + """The boundary arithmetic once each rank's BF16 partial is known (shared with the finalize entry). + + Ranks are summed in FP32 in rank order, the residual added and rounded to BF16 once; 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() + updated = (total + residual.float()).to(torch.bfloat16) + u = 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_103}) + + # 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 (dispatch.MC_ALLREDUCE): the bucket and epilogue options the + # target picks for it (kernels/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): + # dispatch.ATTN_PUSH / dispatch.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) + ), + ) + + 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, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_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, + ) + + 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, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute. + + 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) + + 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, + ) -> 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_103__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..7e6ebeef7d31 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + # 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 (dispatch.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_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_103__tp4/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv.py new file mode 100644 index 000000000000..6a454a7d2af5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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 tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.mx_gemv import op as _gemv + +from ..moe.m3_moe_engine import mxfp8_dequant + +__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_103}) + + # 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 (dispatch.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_103__tp4/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..05304eec0953 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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 tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.mx_gemv import op as _gemv + +from ..moe.m3_moe_engine import mxfp8_quant_dequant +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_103}) + + # 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 (dispatch.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_103__tp4/ops/gemm/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_target_head.py new file mode 100644 index 000000000000..5f3b4ad1e2b0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + # 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_103__tp4/ops/gemm/mxfp8_low_m_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/mxfp8_low_m_gemm.py new file mode 100644 index 000000000000..915c401a3ca1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/mxfp8_low_m_gemm.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused MXFP8 activation quantization + low-M MXFP8 GEMM of the decode projections.""" + +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 + + +def _ue8m0_unswizzle(wsf: torch.Tensor, n: int, k: int) -> torch.Tensor: + """[N, K / 32] scale bytes out of the 128x4-swizzled buffer of ``mxfp8_quantize(w, True)``.""" + cols = k // 32 + padded_cols = (cols + 3) // 4 * 4 + r = torch.arange(n, device=wsf.device).view(-1, 1) + c = torch.arange(cols, device=wsf.device).view(1, -1) + index = ( + (c % 4) + + (c // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) + ) + return wsf.reshape(-1)[index] + + +def _mxfp8_dequant_rows(x: torch.Tensor) -> torch.Tensor: + """The activation as the kernel quantizes it (``trtllm::mxfp8_quantize``), dequantized back to fp32. + + Per 32 columns: scale = the smallest power of two at or above ``amax / 448`` (UE8M0, round up), data = x / scale + rounded to nearest E4M3; an all-zero block stays zero. + """ + m, k = x.shape + blocks = x.float().view(m, k // 32, 32) + amax = blocks.abs().amax(-1, keepdim=True) + 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) + data = torch.where(amax != 0, blocks / scale, torch.zeros_like(blocks)).to(torch.float8_e4m3fn) + return (data.float() * scale).view(m, k) + + +# The target's configuration per (N, K) projection marked for this kernel, and the most rows it routes here +# (dispatch.LOW_M_GEMM_MAX_TOKENS): (warps, cluster_k, row_tiles, max rows). +CONFIGS = { + (2560, 6144): (4, 4, 2, 8), + (6144, 2048): (8, 2, 4, 8), + (6144, 6144): (4, 4, 1, 1), + (6144, 3072): (4, 2, 2, 8), +} + + +def _config_cells(named: tuple[Cell, ...]) -> tuple[Cell, ...]: + """A cell for every (projection, rows) of ``CONFIGS`` that ``named`` does not list: certification is per + projection, since each runs its own configuration.""" + listed = {(cell.spec["n"], cell.spec["k"], cell.spec["m"]) for cell in named} + return tuple( + Cell( + why=f"{n} x {k}, {m} rows", + spec=dict(m=m, n=n, k=k, warps=warps, cluster_k=cluster_k, row_tiles=row_tiles), + ) + for (n, k), (warps, cluster_k, row_tiles, max_m) in CONFIGS.items() + for m in range(1, max_m + 1) + if (n, k, m) not in listed + ) + + +class Mxfp8LowMGemm(OpWrapper): + """``out[M, N] = mxfp8(x)[M, K] @ dequant(w, wsf)[N, K]^T`` in BF16, quantizing ``x`` in the kernel. + + ``w`` is E4M3 [N, K] with UE8M0 block-32 scales ``wsf`` in the 128x4 swizzled layout; ``x`` BF16 [M, K] with unit + column stride. ``(warps, cluster_k, row_tiles)`` is the kernel configuration the target's table picks per + weight shape; each CTA owns ``16 * row_tiles`` weight rows and ``K / cluster_k`` columns, split-K partials meet + in the cluster root and are summed in rank order, one BF16 rounding at the end. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The tp4 projections marked for this kernel (the decode GEMV runs them first and hands over when it declines), + # each at the configuration the target's table assigns and up to the row count it routes here. + CELLS: tuple[Cell, ...] = ( + Cell( + why="QKV with the index q / k projections (2560 x 6144), one token", + spec=dict(m=1, n=2560, k=6144, warps=4, cluster_k=4, row_tiles=2), + ), + Cell( + why="QKV at the configured row limit, 8 tokens", + spec=dict(m=8, n=2560, k=6144, warps=4, cluster_k=4, row_tiles=2), + ), + Cell( + why="o_proj (6144 x 2048), 8 warps x 4 row tiles, odd row count", + spec=dict(m=5, n=6144, k=2048, warps=8, cluster_k=2, row_tiles=4), + ), + Cell( + why="dense MLP gate_up (6144 x 6144), configured for one token only", + spec=dict(m=1, n=6144, k=6144, warps=4, cluster_k=4, row_tiles=1), + ), + Cell( + why="dense MLP down (6144 x 3072), 8 tokens", + spec=dict(m=8, n=6144, k=3072, warps=4, cluster_k=2, row_tiles=2), + ), + ) + # Every other (projection, rows) of its configurations. + CELLS = CELLS + _config_cells(CELLS) + + note = """ + Only the four configurations the target's table assigns are built, each for 1..8 rows (one 8-row MMA tile). + `pdl` changes launch overlap only and is driven off. Activation blocks whose amax is below `448 * 2^-127` take a + flush-to-zero path the reference does not model; no decode activation is that small. + """ + + def __call__( + self, + x: torch.Tensor, + w: torch.Tensor, + wsf: torch.Tensor, + warps: int, + cluster_k: int, + row_tiles: int, + pdl: bool, + ) -> torch.Tensor: + return torch.ops.trtllm.minimax_m3_mxfp8_low_m_gemm( + x, w, wsf, warps, cluster_k, row_tiles, pdl + ) + + def reference( + self, + x: torch.Tensor, + w: torch.Tensor, + wsf: torch.Tensor, + warps: int, + cluster_k: int, + row_tiles: int, + pdl: bool, + ) -> torch.Tensor: + """Dequantized fp32 product of the MXFP8 activation and weight, rounded to BF16 once. + + ``warps``, ``cluster_k``, ``row_tiles`` and ``pdl`` only pick the kernel's split; they do not enter the + arithmetic. + """ + n, k = w.shape + w_scale = torch.exp2(_ue8m0_unswizzle(wsf, n, k).float() - 127.0) + w_deq = (w.float().view(n, k // 32, 32) * w_scale.unsqueeze(-1)).view(n, k) + with true_fp32_matmul(): + out = _mxfp8_dequant_rows(x) @ w_deq.t() + return out.to(torch.bfloat16) + + def is_valid( + self, + x: torch.Tensor, + w: torch.Tensor, + wsf: torch.Tensor, + warps: int, + cluster_k: int, + row_tiles: int, + pdl: bool, + ) -> None: + n, k = w.shape + # The kernel bulk-copies 256-byte scale chunks by computed offset; the op checks the dtype only, so a short + # or linear-layout scale buffer is read past its end or in the wrong order. + need = (n + 127) // 128 * 128 * ((k // 32 + 3) // 4 * 4) + assert wsf.is_contiguous() and wsf.numel() >= need, ( + f"wsf must be the contiguous 128x4-swizzled UE8M0 buffer of {need} bytes; got {wsf.numel()}" + ) + + +mxfp8_low_m_gemm = Mxfp8LowMGemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/ops/moe/m3_moe16_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe16_engine.py new file mode 100644 index 000000000000..a643aa0de12c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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 tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.moe.op import M3MoeLayer +from tensorrt_llm._torch.cute_dsl_kernels.minimax_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_103}) + + # 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_103__tp4/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe_engine.py new file mode 100644 index 000000000000..d8167bb53c0b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe_engine.py @@ -0,0 +1,325 @@ +# 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._torch.cute_dsl_kernels.minimax_m3.moe.op import ( + HIDDEN, + INTER, + SHARED_INTER, + TOP_K, + M3MoeLayer, + M3MoeState, + ar_push_desc, + is_supported, +) +from tensorrt_llm.quantization.utils.fp4_utils import ( + get_reorder_rows_for_gated_act_gemm_row_indices, + get_shuffle_matrix_a_row_indices, +) + +__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_103}) + + # 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_103__tp4/ops/moe/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_router_gemv.py new file mode 100644 index 000000000000..37870d0e6cd7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/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_103}) + + # 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_103__tp4/ops/moe/m3_router_gemv_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_router_gemv_fp4.py new file mode 100644 index 000000000000..58f809e82cf6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/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_103}) + + # 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_103__tp4/ops/moe/m3_router_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_router_tc.py new file mode 100644 index 000000000000..e88af51f7fa2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/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_103}) + + # 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 (dispatch.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_103__tp4/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/m3_embedding_firstnorm.py new file mode 100644 index 000000000000..ff7fb16e2266 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/m3_embedding_firstnorm.py @@ -0,0 +1,89 @@ +# 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 tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.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_103}) + + # 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 (dispatch.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_103__tp4/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/ops/spec/m3_capture_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_capture_norm.py new file mode 100644 index 000000000000..131269ae7d07 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + # 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_103__tp4/ops/spec/m3_chunk_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_chunk_norm.py new file mode 100644 index 000000000000..f6adb2207a20 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + # 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_103__tp4/ops/spec/m3_draft_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_draft_input.py new file mode 100644 index 000000000000..75f7a44dfcfb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + # Every pass of the fused draft loop (dispatch.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_103__tp4/ops/spec/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_draft_step.py new file mode 100644 index 000000000000..a4750f4c1238 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_draft_step.py @@ -0,0 +1,367 @@ +# 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 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_103}) + + # Every pass of the fused draft loop (dispatch.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, + ) -> None: + torch.ops.trtllm.minimax_m3_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, + ) + + 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, + ) -> 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, + ) -> 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" + 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_103__tp4/ops/spec/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_sampled_accept.py new file mode 100644 index 000000000000..e904608fec13 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_sampled_accept.py @@ -0,0 +1,269 @@ +# 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_103}) + + # 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, + ) -> None: + torch.ops.trtllm.minimax_m3_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_103__tp4/ops/spec/m3_sampled_accept_ws_words.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_sampled_accept_ws_words.py new file mode 100644 index 000000000000..3e42fd5c50ad --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/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_sampled_accept`` exchanges through for ``world`` ranks: + two call halves of ``[attempts][rows][world * splits][record words]``.""" + + ARCHS = frozenset({Arch.SM_103}) + + 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)), + ) + + 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_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_103__tp4/ops/spec/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_shard_accept.py new file mode 100644 index 000000000000..90bb11ef84b3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + # Every greedy verify step of the target (dispatch.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_103__tp4/ops/spec/m3_slot_scatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_slot_scatter.py new file mode 100644 index 000000000000..470a6372f2b6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + # 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_103__tp4/ops/spec/m3_step_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_step_gather.py new file mode 100644 index 000000000000..abea2cf620e9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + # 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_103__tp4/runtime.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/runtime.py new file mode 100644 index 000000000000..c40314bf2302 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/runtime.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The step-loop hooks this target opts into (``RuntimeHooks``), all measured on GB300 tp4 AgentX decode. + +* ``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. +* 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. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional, Tuple + +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy, RuntimeHooks + +from .kernels import msa_kv_update +from .spec_fold import M3SpecStepFold + +if TYPE_CHECKING: + from .modeling import ModelingV2MiniMaxM3Nvfp4Sm103Tp4 + + +def make_runtime_hooks( + model: "ModelingV2MiniMaxM3Nvfp4Sm103Tp4", +) -> Tuple[RuntimeHooks, Optional[M3SpecStepFold]]: + """This target's hooks, and the speculative step fold they carry (None without speculation).""" + spec_fold = M3SpecStepFold() if model.spec_config is not None else None + hooks = RuntimeHooks( + pack_step_inputs=True, + skip_unchanged_h2d=True, + loop_on_execution_stream=True, + merge_tuned_tactics_across_ranks=True, + skip_greedy_rng_upload=True, + spec_step_fold=spec_fold, + msa=MsaRuntimePolicy( + short_sequence_identity=True, + kv_update=msa_kv_update.msa_kv_update, + skip_unchanged_block_offsets=True, + reuse_pool_views=True, + ), + ) + return hooks, spec_fold diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec_fold.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec_fold.py new file mode 100644 index 000000000000..593a340a2782 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec_fold.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""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 never stores the sampler's outputs itself, so ``begin_step`` and +``sampler_rows_stored`` always report that nothing was staged. + +* ``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``). +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from .kernels import m3_spec_tail + + +class M3SpecStepFold: + """``SpecStepFold`` implementation of the MiniMax-M3 tp4 target (see the module docstring). + + One instance per target model; the target passes it to ``RuntimeHooks.spec_step_fold``. + """ + + def __init__(self) -> None: + self.store: Optional[object] = None + self.step_inputs: Optional[ + tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] + ] = 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) + + def bind_sampler_store(self, store: object) -> None: + self.store = store + + def begin_step(self, seq_slots: Optional[list[int]], num_rows: int) -> bool: + return False + + 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: + return False + + 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, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec_worker.py new file mode 100644 index 000000000000..773fdb9aa798 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec_worker.py @@ -0,0 +1,1166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Eagle3 one-model speculative decoding of the MiniMax-M3 tp4 target, at stock numerics. + +``M3Eagle3Worker`` is the in-tree ``Eagle3OneModelWorker`` with a fused verify and a fused draft loop: + +* Verify: on decode-only batches whose verify rows the fused kernels cover, the target's BF16 lm_head computes this + rank's vocabulary shard of the logits (the in-tree LM head's own local GEMM; at most 8 verify rows stream a + lossless 12-bit coding of the same BF16 weight, ``dispatch.TARGET_HEAD``) and one kernel replaces the logits + all-gather, the argmax (or the temperature + top-p draw) and the strict acceptance. +* Draft: the in-tree Eagle3 drafter's modules and weights at their precision (BF16 weights and activations, the + full-vocabulary BF16 head, greedy drafts). The first pass may take the per-layer fc_norm RMSNorms in one launch (or + folded into the hidden-state capture) and run its fc beside the target head GEMM (at most 8 verify rows: forked + right after the last captured target layer, ``dispatch.FC_EARLY_MAX_ROWS``). On a decode-only step of at most + 8 requests (``dispatch.DRAFT_STEP``) the draft loop is fused (``_m3_fused_draft_loop``): each pass's draft input in + one launch, the drafter layer's in-tree 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. + +The target's outer CausalLM drives the worker through ``M3SpecDecForwardMixin.forward`` (or the same calls made from +its own forward), and its layer loop captures the Eagle3 hidden states through ``capture_hidden_states``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Optional + +import torch +import triton +import triton.language as tl +from torch import nn + +from tensorrt_llm._torch.distributed import AllReduce, AllReduceParams +from tensorrt_llm._torch.models.modeling_speculative import _slice_spec_position_ids +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager +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, +) +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping + +from . import dispatch +from .kernels import m3_draft_step, m3_sampled_accept, m3_shard_accept, m3_spec_tail, m3_target_head + +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.llmapi.llm_args import EagleDecodingConfig + +# Id of the target's final hidden state in ``layers_to_capture`` (SpecDecOneEngineForCausalLM.layer_idx). +FINAL_HIDDEN_LAYER_ID = -1 + + +@triton.jit +def _m3_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) + + +def draft_ar_norm( + all_reduce: AllReduce, x: torch.Tensor, residual: torch.Tensor, norm: nn.Module +) -> tuple[torch.Tensor, torch.Tensor]: + """One draft-layer boundary in stock Eagle3's structure: the projection's plain all-reduce (BF16 out), then the + norm module with the residual (FlashInfer's fused_add_rmsnorm: FP32 add, norm of the FP32 sum). Returns the + normalized rows and the residual sum.""" + return norm(all_reduce(x), residual) + + +def capture_hidden_states( + spec_metadata: Any, + layer_id: int, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor], + capture_norm: Optional[tuple[torch.Tensor, float]], +) -> 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) (``M3Eagle3Worker.capture_norm``), 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 " + f"buffer capacity: num_tokens={num_tokens}, " + f"capacity={spec_metadata.hidden_states.shape[0]}" + ) + if num_tokens > hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds " + f"available hidden states: num_tokens={num_tokens}, " + f"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 + + +class M3Eagle3Worker(Eagle3OneModelWorker): + """``Eagle3OneModelWorker`` with the MiniMax-M3 fused verify (see the module docstring). + + Methods whose docstring says "Mirrors" are copies of the named in-tree method with the fused-verify branches + inserted; compare them with that method when it changes. + """ + + def __init__( + self, + spec_config: "EagleDecodingConfig", + mapping: Optional[Mapping] = None, + model_config: Optional["ModelConfig"] = None, + use_separate_draft_kv_cache: bool = False, + ): + super().__init__( + spec_config, + mapping=mapping, + model_config=model_config, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + ) + self._init_fused_verify() + + @classmethod + def adopt(cls, worker: Eagle3OneModelWorker) -> "M3Eagle3Worker": + """Turn the in-tree worker the one-engine shell built (``get_spec_worker``) into this worker, in place.""" + if type(worker) is not Eagle3OneModelWorker: + raise TypeError(f"expected an Eagle3OneModelWorker, got {type(worker).__name__}") + worker.__class__ = cls + worker._init_fused_verify() + return worker + + def _init_fused_verify(self) -> None: + # 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_norm[0].weight, stacked fc_norm weights [layers, H]), built eagerly on the first fused first pass. + self._m3_fc_norm_w: Optional[tuple[torch.Tensor, torch.Tensor]] = None + # Armed per step by arm_capture_norm: (stacked fc_norm weights, eps) or None. + self._m3_capture_norm: Optional[tuple[torch.Tensor, float]] = 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 + # (capture_hidden_states), and the fork that did it, joined by _m3_target_accept. + self._m3_fc_early_drafter: Optional[nn.Module] = 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 at <= 8 rows, built before the + # executor's warmup (prebuild_target_head). + self._m3_target_head: Optional[m3_target_head.TargetHead] = None + self._m3_force_vec_warm = False + # The fragment-ordered copy of the drafter's BF16 head shard the fused draft loop streams, built before the + # executor's warmup (prebuild_draft_head). + self._m3_draft_head: Optional[m3_draft_step.DraftHead] = 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 + + @property + def capture_norm(self) -> Optional[tuple[torch.Tensor, float]]: + """What ``capture_hidden_states`` needs for this step's captures (set by ``arm_capture_norm``).""" + return self._m3_capture_norm + + def capture_hidden_states( + self, + spec_metadata: Any, + layer_id: int, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor] = None, + ) -> None: + """``capture_hidden_states`` with this step's capture norm; 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``).""" + capture_hidden_states( + 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 _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 + ) + ) + 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 because + # the accepted token count affects which Mamba states are valid. The + # isinstance check below naturally no-ops on non-Mamba kv_cache_managers, + # so this is safe to run unconditionally regardless of spec mode (Eagle3 + # over a Mamba-style draft is plausible, even if no such draft exists today). + 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: torch.LongTensor, + position_ids: torch.LongTensor, + hidden_states: torch.Tensor, + accepted_tokens: torch.Tensor, + attn_metadata: Any, + spec_metadata: Any, + draft_model: nn.Module, + ): + """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: Any, spec_metadata: Any) -> 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) > dispatch.DECODE_MAX_TOKENS: + return False + return not spec_metadata.wants_advanced_draft_sampling + + def tail_eligible(self, attn_metadata: Any, spec_metadata: Any, 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 ``dispatch.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 self._m3_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 dispatch.SHARD_ACCEPT.covers(key): + dispatch.uncovered( + dispatch.SHARD_ACCEPT.path, + f"draft length {draft_len} 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: Any, spec_metadata: Any, draft_model: Optional[nn.Module] + ) -> 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.""" + return self._m3_structural_ok(attn_metadata, spec_metadata) and self._m3_fc_norm_ok( + draft_model + ) + + def _m3_seed_first_pass( + self, num_accepted_tokens: torch.Tensor, spec_metadata: Any, 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: torch.Tensor, + accepted_tokens: torch.Tensor, + num_tokens: int, + num_contexts: int, + spec_metadata: Any, + draft_model: nn.Module, + ) -> Optional[dict[str, Any]]: + """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 in-tree 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 not self._m3_fc_norm_ok(draft_model) + ): + return None + dm = draft_model.model + layers = len(spec_metadata.layers_to_capture) + if dm.hidden_size_in != spec_metadata.hidden_size or layers != len(dm.fc_norm): + return None + cache = self._m3_fc_norm_w + if cache is None or cache[0] is not dm.fc_norm[0].weight: + if torch.cuda.is_current_stream_capturing(): + dispatch.capture_missed("spec_tail", None, "the stacked Eagle3 fc_norm weights") + return None + cache = ( + dm.fc_norm[0].weight, + torch.stack([n.weight.detach() for n in dm.fc_norm]).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], dm.fc_norm[0].variance_epsilon) + hidden = dm.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} + + @staticmethod + def _m3_fc_norm_ok(draft_model: Optional[nn.Module]) -> bool: + """An Eagle3 drafter with an fc over several captured layers, each normalized by its own fc_norm (one + shared epsilon): what the fused first drafter inputs compute.""" + dm = getattr(draft_model, "model", None) + fc_norm = getattr(dm, "fc_norm", None) + return ( + fc_norm is not None + and getattr(dm, "fc", None) is not None + and len({n.variance_epsilon for n in fc_norm}) == 1 + ) + + def _m3_fc_fork( + self, + spec_metadata: Any, + num_tokens: int, + draft_model: Optional[nn.Module], + 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 or not self._m3_fc_norm_ok(draft_model): + return None + dm = draft_model.model + layers = len(spec_metadata.layers_to_capture) + if ( + dm.hidden_size_in != spec_metadata.hidden_size + or layers != len(dm.fc_norm) + 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], dm.fc_norm[0].variance_epsilon, use_pdl=False + ) + self._m3_fc_hidden = dm.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 or + torch.compile, 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: Any, + spec_metadata: Any, + lm_head: nn.Module, + draft_model: Optional[nn.Module], + ) -> 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 + ``dispatch.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 + # 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) + <= dispatch.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 self._m3_fc_norm_ok(draft_model) + and fc_norm_w[0] is draft_model.model.fc_norm[0].weight + and len(spec_metadata.layers_to_capture) == fc_norm_w[1].shape[0] + ): + self._m3_capture_norm = (fc_norm_w[1], draft_model.model.fc_norm[0].variance_epsilon) + + @staticmethod + def _m3_sampled_verify_ok(spec_metadata: Any) -> 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: torch.Tensor, + spec_metadata: Any, + batch_size: int, + draft_len: int, + draft_model: Optional[nn.Module] = None, + ): + """Fused verification (see tail_eligible); same outputs as the strict acceptance path.""" + 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_kernel[(1,)]( + 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=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 + ) + 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 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 in-tree LM 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: + logits = m3_target_head.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.shard_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + force_total, + self._m3_next_new, + force_vec=force_vec, + ) + else: + n_rows = rows.shape[0] + seeds, offsets = self._rng_state_per_token(spec_metadata, n_rows) + result = m3_sampled_accept.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, + ) + 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: + raise RuntimeError( + "M3 fused verify unavailable after tail_eligible() accepted the step" + ) + 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 <= 8 rows (rebuilt from + the current weight values on every call). Called by the target's ``cache_derived_state`` once the weights are + real, before the executor's warmup, 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 in-tree GEMM runs.""" + self._m3_target_head = None + if not self._m3_sharded_head(lm_head) or not m3_target_head.head_supported(lm_head): + return False + head = m3_target_head.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 in-tree BF16 head GEMM" + ) + return head.coded + + def _m3_coded_head(self, lm_head: nn.Module, rows: int) -> Optional[m3_target_head.TargetHead]: + """The coded target head for a fused verify of ``rows`` rows, or None for the in-tree GEMM: above + ``dispatch.TARGET_HEAD_MAX_ROWS`` rows (uncovered), on a shard whose coding failed its round trip + (uncovered: the in-tree GEMM is the exact path), and when the copy was not built for this weight (a miss).""" + path = dispatch.TARGET_HEAD.path + if not dispatch.TARGET_HEAD.covers(rows): + dispatch.uncovered(path, f"{rows} verify rows (the in-tree head GEMM runs)") + return None + head = self._m3_target_head + if head is None or head.weight is not lm_head.weight: + dispatch.missed( + path, + rows, + "the coded target head is not built for the head's current weight (prebuild_target_head)", + ) + return None + if not head.coded: + dispatch.uncovered( + path, "the coded shard failed its bitwise round trip (the in-tree head GEMM runs)" + ) + return None + return head + + def _m3_force_vec( + self, batch_size: int, base_total: int, max_total: int, device: torch.device + ) -> Optional[torch.Tensor]: + """Per-request forced accepted counts for a fractional forced acceptance, else None. + + Draws exactly as _apply_force_accepted_tokens does for a decode-only batch (same fixed-seed pool, one counter + bump per step, the same (counter, slot) hash), so every TP rank and the generic path agree on the counts. + """ + 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) + # One launch: bump the shared counter and draw, with _apply_force_accepted_tokens' exact arithmetic. + _m3_force_vec_kernel[(1,)]( + 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=max(16, 1 << (batch_size - 1).bit_length()), + ) + return out + + @staticmethod + def _m3_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 + ) + + # ---------------------------------------------------------------------------------------------- fused draft loop + + def prebuild_draft_head(self, draft_model: Optional[nn.Module]) -> bool: + """Build the fragment-ordered copy of the drafter's BF16 head shard that the fused draft loop streams (rebuilt + from the current weight values on every call). Called by the target's ``cache_derived_state`` once the weights + are real, before the executor's warmup, 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.prepare_head(draft_model.lm_head, self.mapping.tp_rank) + return True + + def _m3_drafter_supported(self, draft_model: nn.Module) -> bool: + """The drafter the fused draft loop runs: one plain (non-MLA) Eagle3 layer with a gated MLP over the + target-vocabulary embedding (whole table on every rank), a draft vocabulary equal to the target's (no d2t), + all-reduced TP projections, and a bias-free vocab-sharded BF16 head without padding.""" + dm = getattr(draft_model, "model", None) + if ( + dm is None + or getattr(dm, "num_layers", 0) != 1 + or dm._use_mla + or dm._eh_proj_before_attn + ): + return False + layer = dm.midlayer + if layer._next_layer_regular or not isinstance(layer.mlp, GatedMLP): + return False + embed = dm.embed_tokens + if embed is None or getattr(embed, "tp_size", 1) != 1: + return False + if self.mapping is None or self.mapping.tp_size <= 1 or self.mapping.enable_attention_dp: + return False + if getattr(layer.self_attn.o_proj, "all_reduce", None) is None: + return False + if getattr(layer.mlp.down_proj, "all_reduce", None) is None: + return False + if self._d2t is not None or getattr(dm, "d2t", None) is not None: + return False + lm_head = draft_model.lm_head + return self._m3_sharded_head(lm_head) and m3_draft_step.head_supported(lm_head) + + def _m3_fused_loop_ok( + self, + attn_metadata: Any, + spec_metadata: Any, + batch_size: int, + draft_model: Optional[nn.Module], + ) -> 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) + ``dispatch.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 = dispatch.DRAFT_STEP.path + if not self._m3_drafter_supported(draft_model) or not hasattr( + attn_metadata, "kv_lens_cuda" + ): + dispatch.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 dispatch.DRAFT_STEP.covers((draft_len, batch_size)): + dispatch.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: + dispatch.missed( + path, + batch_size, + "the drafter head's fragment-ordered copy is not built for its current weight (prebuild_draft_head)", + ) + return False + return m3_draft_step.workspace(self.mapping, head.wf.device) is not None + + 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: dict[str, Any], + attn_metadata: Any, + spec_metadata: Any, + draft_model: nn.Module, + draft_kv_cache_manager: Any, + batch_size: int, + num_accepted_tokens: torch.Tensor, + ) -> torch.Tensor: + """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_layernorm 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 in-tree attention and MLP, each followed by the in-tree all-reduce with the residual add + and the next RMSNorm fused into it (post_attention_layernorm, 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 + dm = draft_model.model + layer = dm.midlayer + attn = layer.self_attn + mlp = layer.mlp + attn_all_reduce = attn.o_proj.all_reduce + mlp_all_reduce = mlp.down_proj.all_reduce + table = dm.embed_tokens.weight + 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:] + + no_all_reduce = AllReduceParams(enable_allreduce=False) + 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, + layer.input_layernorm.weight, + layer.input_layernorm.variance_epsilon, + layer.hidden_norm.weight, + layer.hidden_norm.variance_epsilon, + ) + h = attn( + position_ids=pos, + hidden_states=x, + attn_metadata=md, + all_reduce_params=no_all_reduce, + ) + h, residual = draft_ar_norm( + attn_all_reduce, h, residual, layer.post_attention_layernorm + ) + h = mlp(h, final_all_reduce_params=no_all_reduce) + h, to_save = draft_ar_norm(mlp_all_reduce, h, residual, dm.norm) + if dm._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.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 + + +class M3SpecDecForwardMixin: + """The target CausalLM's step forward with the fused-verify hooks of the speculative wrapper. + + Mirrors ``SpecDecOneEngineForCausalLM.forward`` (tensorrt_llm/_torch/models/modeling_speculative.py) with the + fused-verify branches: before the target forward the worker arms the capture norm, and a step the worker verifies + from the logits shard (``M3Eagle3Worker.tail_eligible``) skips the gathered target logits. + + The class it is mixed into (listed before ``DecoderModelForCausalLM``) provides ``model`` (called like the in-tree + decoder model, capturing through ``M3Eagle3Worker.capture_hidden_states``), ``lm_head``, ``logits_processor``, + ``model_config``, ``draft_model`` and ``spec_worker`` (an ``M3Eagle3Worker``, or None without speculation). + """ + + def forward( + self, + attn_metadata: Any, + 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[Any] = None, + resource_manager: Optional[Any] = None, + **kwargs, + ) -> Any: + spec_worker: Optional[M3Eagle3Worker] = self.spec_worker + if spec_worker is not None and spec_metadata is not None: + spec_worker.arm_capture_norm( + attn_metadata, spec_metadata, self.lm_head, self.draft_model + ) + hidden_states = self.model( + input_ids=input_ids, + attn_metadata=attn_metadata, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + **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): + if spec_worker is not None: + spec_worker.capture_hidden_states( + spec_metadata, FINAL_HIDDEN_LAYER_ID, hidden_states, None + ) + else: + capture_hidden_states( + spec_metadata, FINAL_HIDDEN_LAYER_ID, hidden_states, None, None + ) + if attn_metadata.padded_num_tokens is not None: + hidden_states = hidden_states[: attn_metadata.num_tokens] + + if spec_worker is not None: + # The fused verify computes this rank's logits shard itself and needs no gathered logits. + if spec_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 + ) + return spec_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, + ) + + return self.logits_processor.forward( + hidden_states, + self.lm_head, + attn_metadata, + return_context_logits, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/weights.py new file mode 100644 index 000000000000..03d9b08caa6b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/weights.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight loading: MiniMax-M3 NVFP4 / sm_103 / tp4. + +The checkpoint is the vision-language wrapper's. ``language_model.*`` holds the +language model; ``vision_tower.*``, ``multi_modal_projector.*`` and +``patch_merge_mlp.*`` hold the vision tower and its projector. This target is +text only, so the keys split three ways: + +* ``language_model.*`` goes to the built-in MiniMax-M3 text loader with the + prefix stripped (its multi-token-prediction head, ``model.mtp.*``, is not + loaded, as in the built-in model). That loader packs the five-way + Q/K/V/index-Q/index-K projection, folds the Gemma boundary norms and shards + every tensor for tp4; the modules it fills are this target's subclasses of + the built-in ones, under the built-in names. +* The vision tower and the projector are a predicted non-load: listed here, + never read. +* Any other key fails the load, naming it, rather than being dropped. +""" + +from typing import Any, Dict, Optional + +from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper +from tensorrt_llm._torch.models.modeling_minimaxm3 import ( + MiniMaxM3ForCausalLM, + _strip_language_model_prefix, +) + +_LANG_PREFIX = "language_model." + +# The vision tower's and the projector's key families: in the checkpoint, not loaded by this text-only target. +PREDICTED_NON_LOAD = ("vision_tower.", "multi_modal_projector.", "patch_merge_mlp.") + + +def load( + model, + weights: Dict[str, Any], + weight_mapper: Optional[BaseWeightMapper] = None, + params_map: Optional[Dict[str, str]] = None, + allow_partial_loading: bool = False, +) -> None: + """Load the language model's weights into ``model``, the target shell, and check every other key is predicted.""" + unknown = sorted( + k + for k in weights.keys() + if not k.startswith(_LANG_PREFIX) and not k.startswith(PREDICTED_NON_LOAD) + ) + assert not unknown, ( + f"{len(unknown)} checkpoint key(s) are neither language-model weights nor a predicted non-load, " + f"first {unknown[:5]}" + ) + text_weights, _ = _strip_language_model_prefix( + {k: v for k, v in weights.items() if k.startswith(_LANG_PREFIX)} + ) + assert text_weights, f"the checkpoint has no {_LANG_PREFIX}* keys" + MiniMaxM3ForCausalLM.load_weights( + model, + text_weights, + weight_mapper=weight_mapper, + params_map=params_map, + allow_partial_loading=allow_partial_loading, + ) 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 new file mode 100644 index 000000000000..480ea23f9b02 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/routing.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Where a MiniMaxM3SparseForConditionalGeneration config lands. Read this file and you know. + +``route`` is one forward-reading decision: the facts that fix which code runs, +compared against the deployments the targets below were written for. When none +fits it raises ``NoTarget`` stating those facts. In ``auto`` the engine then +uses the built-in MiniMax-M3 implementation; in ``require`` it raises that error. + +The checkpoint declares the vision-language wrapper, and the language model's +shape lives in its ``text_config``, so that is what this tree reads. The +targets are text-only: they load no vision tower and refuse image input. +""" + +from __future__ import annotations + +from typing import Any, Optional, Tuple + +from tensorrt_llm.llmapi.llm_args import EagleDecodingConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +from ..._router_index import NoTarget + +# 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. +_SM = (10, 3) + +# 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. +# +# (num_hidden_layers, hidden_size, num_local_experts, intermediate_size, +# num_attention_heads, num_key_value_heads). The expert width and the head +# counts fix the per-rank operand shapes the decode kernels are compiled for. +_SHAPE = (60, 6144, 128, 3072, 64, 4) + + +def _field(config: Any, name: str) -> Any: + """Read ``name`` off a sub-config that may be an object or a plain dict.""" + if isinstance(config, dict): + return config.get(name) + return getattr(config, name, None) + + +def _quant_algo(pretrained_config: Any) -> Optional[QuantAlgo]: + """The quantization the checkpoint's config declares, as the engine reads it. + + The NVFP4 checkpoint is a modelopt mixed-precision export (NVFP4 routed + experts, MXFP8 dense projections, bf16 router, embedding and head), declared + in ``config.json`` as ``MIXED_PRECISION``; the per-layer split lives in + ``hf_quant_config.json``, which ``route`` does not see. + """ + from tensorrt_llm._torch.model_config import ModelConfig + + declared = getattr(pretrained_config, "quantization_config", None) + if declared is None: + return None + try: + quant_config, _ = ModelConfig.load_hf_quant_config(declared, "TRTLLM") + except ValueError: + return None + return quant_config.quant_algo + + +def _parallel(m) -> Optional[str]: + """Name the parallel topology, or None if no target implements it. + + 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), + so it selects a target rather than a runtime branch. + """ + if not ( + m.world_size == 4 + and m.tp_size == 4 + and m.pp_size == 1 + and m.cp_size == 1 + and m.moe_tp_size == 4 + and m.moe_ep_size == 1 + and not m.enable_attention_dp + ): + return None + return "tp4" + + +def _eagle3_one_model_or_none(spec_config: Optional[Any]) -> bool: + """No speculation, or the one-model Eagle3 loop the target's worker extends.""" + if spec_config is None: + return True + return ( + isinstance(spec_config, EagleDecodingConfig) + and spec_config.spec_dec_mode.is_eagle3_one_model() + ) + + +def route( + pretrained_config: Any, + mapping: Any, + compute_capability: Tuple[int, int], + spec_config: Optional[Any], +) -> type: + """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. + + Each target is imported where it is returned, so declining costs no import + of modeling code. + """ + text_config = getattr(pretrained_config, "text_config", None) + shape = ( + None + if text_config is None + else tuple( + _field(text_config, name) + for name in ( + "num_hidden_layers", + "hidden_size", + "num_local_experts", + "intermediate_size", + "num_attention_heads", + "num_key_value_heads", + ) + ) + ) + quant_algo = _quant_algo(pretrained_config) + parallel = _parallel(mapping) + speculation = _eagle3_one_model_or_none(spec_config) + + # MiniMax-M3 NVFP4 on 4 GB300s (sm103), tensor-parallel attention and + # experts, one-model Eagle3 or no speculative decoding. + if ( + compute_capability == _SM + and shape == _SHAPE + and quant_algo == QuantAlgo.MIXED_PRECISION + and parallel == "tp4" + and speculation + ): + from .minimax_m3_nvfp4__sm_103__tp4.modeling import ModelingV2MiniMaxM3Nvfp4Sm103Tp4 + + return ModelingV2MiniMaxM3Nvfp4Sm103Tp4 + + raise NoTarget( + f"compute capability {compute_capability}, text_config shape {shape}, " + f"quantization {quant_algo.name if quant_algo is not None else 'none'}, parallel " + f"ws={mapping.world_size} tp={mapping.tp_size} pp={mapping.pp_size} cp={mapping.cp_size} " + f"moe_tp={mapping.moe_tp_size} moe_ep={mapping.moe_ep_size} " + f"attention_dp={mapping.enable_attention_dp}, speculative decoding " + f"{spec_config.spec_dec_mode.name if spec_config is not None else 'off'}; the one " + "minimax_m3 target is the MiniMax-M3 NVFP4 checkpoint on sm103 (GB300), 4 ranks of " + "tensor-parallel attention and experts, with one-model Eagle3 or no speculative decoding" + ) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/msa_decode.py b/tensorrt_llm/_torch/attention/backends/fmha/msa_decode.py index c6c6af4d08cd..b4b41ad3a645 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/msa_decode.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/msa_decode.py @@ -39,6 +39,21 @@ from ..sparse.minimax_m3.kernels.trtllm_gen_dense_decode import DenseDecodeWorkspaceLayout +def dense_decode_max_seq_len(attn, max_seq_len: int, num_requests: int) -> int: + """The max_seq_len trtllm-gen dense decode is handed for a step of ``num_requests`` requests. + + A model may cap it per request count through its attention module's + ``dense_decode_max_seq_len_by_batch`` ({requests: max_seq_len}). The value + only sizes the kernel's multi-CTA KV split: every CTA still covers its share + of each request's actual length, so a cap changes the result only by the + split-merge order. The request count is fixed per CUDA graph, so the choice + is graph-safe. + """ + caps = getattr(attn, "dense_decode_max_seq_len_by_batch", None) + cap = caps.get(num_requests) if caps else None + return max_seq_len if cap is None else min(max_seq_len, int(cap)) + + class MsaDecodeFmha(PhasedFmha): """MiniMax-M3 generation attention on the Triton and trtllm-gen kernels. @@ -261,7 +276,9 @@ def _run_dense( sm_scale=(head_dim**-0.5) / float(attn.q_scaling), output=params.output.view(num_tokens, attn.num_heads, head_dim), decode_query_len=params.input_seq_length, - max_seq_len=int(metadata.msa_max_kv_len), + max_seq_len=dense_decode_max_seq_len( + attn, int(metadata.msa_max_kv_len), int(seq_lens.shape[0]) + ), max_num_requests=int(metadata.max_num_requests), staged_subpage_table=staged_table, staged_subpages_per_slot=staged_factor, @@ -270,4 +287,4 @@ def _run_dense( ) -__all__ = ["MsaDecodeFmha"] +__all__ = ["MsaDecodeFmha", "dense_decode_max_seq_len"] diff --git a/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py b/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py index 120feaafeeee..bc1301e57c1b 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py @@ -228,24 +228,39 @@ def run_msa_nvfp4_sparse_gqa( ) cu_q = metadata.msa_cu_q_lens[: batch + 1] cu_kv = metadata.msa_cu_kv_lens[: batch + 1] - q2k = kv_block_indexes.permute(1, 0, 2).contiguous() - topk = int(q2k.shape[-1]) + topk = int(kv_block_indexes.shape[-1]) page_size = int(k_paged.shape[2]) if topk != MSA_REQUIRED_TOPK: raise ValueError(f"MiniMax-M3 MSA NVFP4 requires topK={MSA_REQUIRED_TOPK}, got {topk}") - k2q_row_ptr, k2q_q_indices, schedule = sparse.build_k2q_csr( - q2k, - cu_q, - cu_kv, - page_size, - total_k=total_k, - max_seqlen_k=max_kv_len, - max_seqlen_q=max_q_len, - total_rows=total_k_rows, - qhead_per_kv=int(q.shape[1]) // int(k_paged.shape[1]), - return_schedule=True, - ) + # A short pure-prefill step hands every layer the same identity selection (staged once per step by the + # backend), so its k2q CSR worklist is layer-invariant: build it for the first layer and reuse it. + identity = metadata._msa_prefill_identity + csr_key = (batch, int(kv_block_indexes.shape[0])) + cached = metadata._msa_identity_csr + if ( + identity is not None + and cached is not None + and cached[0] is identity + and cached[1] == csr_key + ): + k2q_row_ptr, k2q_q_indices, schedule = cached[2] + else: + q2k = kv_block_indexes.permute(1, 0, 2).contiguous() + k2q_row_ptr, k2q_q_indices, schedule = sparse.build_k2q_csr( + q2k, + cu_q, + cu_kv, + page_size, + total_k=total_k, + max_seqlen_k=max_kv_len, + max_seqlen_q=max_q_len, + total_rows=total_k_rows, + qhead_per_kv=int(q.shape[1]) // int(k_paged.shape[1]), + return_schedule=True, + ) + if identity is not None: + metadata._msa_identity_csr = (identity, csr_key, (k2q_row_ptr, k2q_q_indices, schedule)) result = sparse.sparse_atten_nvfp4_kv_func( q, k_paged.view(torch.uint8), 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 ca94afce9b69..0b0f458224cd 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,7 +28,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple import torch @@ -57,6 +57,37 @@ from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig +class _UnchangedCopy: + """Remembers the inputs of the last decode-only block-table copy into one destination. + + In steady decode the block tables change only when a request grows into a + new page, so a copy whose destination, batch order, decode-only shape and + host block-offset table equal the previous copy's can be skipped. + """ + + def __init__(self) -> None: + self._key: Optional[tuple] = None + self._host: Optional[torch.Tensor] = None + + def skip(self, dst, request_ids, beam_width, num_contexts, num_seqs, max_blocks, host) -> bool: + if num_contexts != 0: + self._key = None + return False + key = (dst.data_ptr(), tuple(request_ids), beam_width, num_seqs, max_blocks) + return key == self._key and self._host is not None and torch.equal(host, self._host) + + def record( + self, dst, request_ids, beam_width, num_contexts, num_seqs, max_blocks, host + ) -> None: + if num_contexts != 0: + return + self._key = (dst.data_ptr(), tuple(request_ids), beam_width, num_seqs, max_blocks) + if self._host is None or self._host.shape != host.shape: + self._host = host.clone() + else: + self._host.copy_(host) + + class MiniMaxM3SparseIndexCache: """Plain-tensor side cache for the M3 sparse index branch. @@ -260,6 +291,10 @@ class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2): own pool. """ + # Decode-step shortcuts a model opts into (RuntimeHooks.msa); see configure_decode_shortcuts. + skip_unchanged_block_offsets = False + reuse_pool_views = False + _main_kv_mapper_kind = MapperKind.NHD # Virtual attention-op pools for shared draft layers: # (attention_op_pool_id, source_storage_pool_id) plus their copy parameters. @@ -330,6 +365,8 @@ def __init__( self._kv_buffers: dict[tuple[int, str], torch.Tensor] = {} self._block_scale_buffers: dict[tuple[int, str], torch.Tensor] = {} super().__init__(*args, **kwargs) + self._unchanged_copies: Dict[int, _UnchangedCopy] = {} + self._pool_views: Dict[tuple, torch.Tensor] = {} if self.dtype == DataType.NVFP4: dense_target_layers = [ @@ -586,6 +623,20 @@ def copy_batch_block_offsets( ) -> None: if self.num_attention_op_pools == 0: return + if self.skip_unchanged_block_offsets: + guard = self._unchanged_copies.setdefault(dst_tensor.data_ptr(), _UnchangedCopy()) + args = ( + dst_tensor, + request_ids, + beam_width, + num_contexts, + num_seqs, + max_blocks, + self.host_kv_cache_block_offsets, + ) + if guard.skip(*args): + return + guard.record(*args) super().copy_batch_block_offsets( dst_tensor, request_ids, beam_width, num_contexts, num_seqs, max_blocks=max_blocks ) @@ -769,9 +820,56 @@ def get_buffers( addr_key, torch_dtype, num_slots, scale, page_shape = self._kv_slot_geometry( layer_idx, kv_layout ) - full_slot_shape = [num_slots, scale, *page_shape] - full_view = convert_to_torch_tensor(TensorWrapper(addr_key, torch_dtype, full_slot_shape)) - return full_view[:, :2] + return self._pool_view("kv", addr_key, torch_dtype, num_slots, scale, page_shape) + + def configure_decode_shortcuts( + self, skip_unchanged_block_offsets: bool, reuse_pool_views: bool + ) -> None: + """Turn on the decode-step shortcuts a model opted into. + + ``skip_unchanged_block_offsets``: skip a decode-only block-table copy + 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. + """ + self.skip_unchanged_block_offsets = skip_unchanged_block_offsets + self.reuse_pool_views = reuse_pool_views + + def _pool_view( + self, + kind: str, + addr_key: int, + torch_dtype: torch.dtype, + num_slots: int, + scale: int, + page_shape: List[int], + ) -> torch.Tensor: + """The ``[num_slots, 2, ...]`` K+V view of a pool rooted at ``addr_key``. + + The view is a pure function of the pool geometry, which is static once the + pool is laid out, so with ``reuse_pool_views`` it is built once per geometry + (a re-laid-out pool has another key). Inference-mode views are inference + tensors, so each mode keeps its own. + """ + key = ( + kind, + addr_key, + torch_dtype, + num_slots, + scale, + tuple(page_shape), + torch.is_inference_mode_enabled(), + ) + view = self._pool_views.get(key) if self.reuse_pool_views else None + if view is None: + full_slot_shape = [num_slots, scale, *page_shape] + full_view = convert_to_torch_tensor( + TensorWrapper(addr_key, torch_dtype, full_slot_shape) + ) + view = full_view[:, :2] + if self.reuse_pool_views: + self._pool_views[key] = view + return view def _kv_scale_slot_geometry( self, layer_idx: int, kv_layout: Optional[str] @@ -842,9 +940,7 @@ def get_block_scale_buffers( addr_key, torch_dtype, num_slots, scale, page_shape = self._kv_scale_slot_geometry( layer_idx, kv_layout ) - full_slot_shape = [num_slots, scale, *page_shape] - full_view = convert_to_torch_tensor(TensorWrapper(addr_key, torch_dtype, full_slot_shape)) - return full_view[:, :2] + return self._pool_view("scale", addr_key, torch_dtype, num_slots, scale, page_shape) def get_kv_subpage_pool( self, layer_idx: int, kv_layout: str = "HND" @@ -999,6 +1095,7 @@ def __init__(self, manager: MiniMaxM3KVCacheManagerV2, draft_layer_ids: Sequence f"got {len(draft_layer_ids)}. Multiple layers require separate pool roots." ) self._manager = manager + self._unchanged_copy = _UnchangedCopy() layer_id = draft_layer_ids[0] local = int(manager.layer_offsets[layer_id]) self._local_layer_idx = local @@ -1063,6 +1160,19 @@ def copy_batch_block_offsets( """ assert beam_width == 1, "beam_width must be 1 for KVCacheManagerV2" manager = self._manager + if manager.skip_unchanged_block_offsets: + args = ( + dst_tensor, + request_ids, + beam_width, + num_contexts, + num_seqs, + max_blocks, + self.host_kv_cache_block_offsets, + ) + if self._unchanged_copy.skip(*args): + return + self._unchanged_copy.record(*args) copy_idx = manager.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) assert copy_idx.shape[0] == num_seqs copy_batch_block_offsets_to_device( diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py index f288495eeff1..d08488e68e24 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py @@ -204,8 +204,8 @@ def from_sparse_params( class PagedKvSlotMapping(NamedTuple): """One step's paged-cache slot mapping (see build_paged_kv_slot_mapping).""" - req_to_token: torch.Tensor - slot_ids: torch.Tensor + req_to_token: Optional[torch.Tensor] + slot_ids: Optional[torch.Tensor] out_cache_loc: torch.Tensor block_ids_cpu: torch.Tensor @@ -217,6 +217,7 @@ def build_paged_kv_slot_mapping( qo_lens_cpu: torch.Tensor, qo_offset_cpu: torch.Tensor, device: torch.device, + device_tables: bool = True, ) -> PagedKvSlotMapping: """Build the backend-neutral paged-cache slot mapping. @@ -239,6 +240,9 @@ def build_paged_kv_slot_mapping( staged with non-blocking copies, so this neither syncs the host nor reads device memory back. It still allocates per step, so call it from prepare(), never from the forward path. + + With device_tables=False, req_to_token and slot_ids are not built (None) and out_cache_loc stays a pinned host + tensor, for callers that only stage out_cache_loc into their own persistent buffer. """ tokens_per_block = int(kv_cache_manager.tokens_per_block) # block_ids_per_seq is a [batch, max_blocks_per_seq] tensor; row b holds the @@ -250,13 +254,15 @@ def build_paged_kv_slot_mapping( max_blocks = int(block_ids.shape[1]) max_kv_len = max_blocks * tokens_per_block - # Expand block ids -> per-token slot ids. - block_ids_dev = block_ids.to(device, non_blocking=True).to(torch.int64) - within_block = torch.arange(tokens_per_block, device=device, dtype=torch.int64) - # Outer product per batch entry: [batch, max_blocks, tokens_per_block] - slot_grid = block_ids_dev.unsqueeze(-1) * tokens_per_block + within_block - req_to_token = slot_grid.reshape(batch, max_kv_len).to(torch.int32) - slot_ids = torch.arange(batch, device=device, dtype=torch.int32) + req_to_token = slot_ids = None + if device_tables: + # Expand block ids -> per-token slot ids. + block_ids_dev = block_ids.to(device, non_blocking=True).to(torch.int64) + within_block = torch.arange(tokens_per_block, device=device, dtype=torch.int64) + # Outer product per batch entry: [batch, max_blocks, tokens_per_block] + slot_grid = block_ids_dev.unsqueeze(-1) * tokens_per_block + within_block + req_to_token = slot_grid.reshape(batch, max_kv_len).to(torch.int32) + slot_ids = torch.arange(batch, device=device, dtype=torch.int32) # out_cache_loc: per-new-token slot ids, in flattened query-token order. # Expanding the per-request lengths on the host reproduces the same slot @@ -281,7 +287,10 @@ def build_paged_kv_slot_mapping( block_ids[row, block_col].to(torch.long) * tokens_per_block + (pos - block_col * tokens_per_block) ).to(torch.int32) - out_cache_loc = async_tensor_h2d(out_cache_loc_cpu, torch.int32, device) + if device_tables: + out_cache_loc = async_tensor_h2d(out_cache_loc_cpu, torch.int32, device) + else: + out_cache_loc = maybe_pin_memory(out_cache_loc_cpu) return PagedKvSlotMapping(req_to_token, slot_ids, out_cache_loc, block_ids) 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 eb5443442153..dd6d0b5c7436 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 @@ -38,6 +38,9 @@ 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.runtime_hooks import MsaRuntimePolicy from tensorrt_llm._utils import maybe_pin_memory from tensorrt_llm.bindings import DataType from tensorrt_llm.models.modeling_utils import QuantConfig @@ -102,6 +105,17 @@ def _worst_case_proxy_max_k_tiles( return int(proxy_plan[3]["max_k_tiles"]) +def _host_prefix_sum(lens_cpu: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """[0, cumsum(lens)...] as a pinned host tensor of ``dtype``.""" + sums = torch.cumsum(lens_cpu.to(torch.int64), 0) + return maybe_pin_memory(torch.cat([sums.new_zeros(1), sums]).to(dtype)) + + +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) + + class MsaDecodeSpan(NamedTuple): """Generation rows served by MiniMax-M3's dedicated decode kernels. @@ -160,6 +174,13 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): msa_kv_indices: Optional[torch.Tensor] = None msa_max_score: Optional[torch.Tensor] = None msa_n_valid_blocks: Optional[torch.Tensor] = None + # Decode-step preparation the model opted into (RuntimeHooks.msa). + msa_policy: MsaRuntimePolicy = MsaRuntimePolicy() + # Short-sequence steps (MsaRuntimePolicy.short_sequence_identity): every query token has at most topk valid + # blocks, so the selection is the identity table [kv_heads, tokens, topk] = (i < n_valid ? i : -1), staged here + # by prepare(). _msa_short_ok says whether this step staged it. + msa_short_table: Optional[torch.Tensor] = None + _msa_short_ok: bool = False # The same page table and lengths as msa_kv_indices / msa_kv_lens, in the # per-request 2-D form the decode kernels index directly # (block_table[request, block] and seq_lens[request]). fmha_sm100 instead @@ -203,6 +224,12 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): # _msa_prefill_n_valid_buf is the persistent backing store for the view. _msa_prefill_n_valid_buf: Optional[torch.Tensor] = None _msa_prefill_n_valid_blocks: Optional[torch.Tensor] = None + # Pure-prefill step where every query token has at most topk valid blocks: selection is the identity, so the + # indexer returns this [total_q, topk] table (per KV head) instead of scoring (built once per step). + _msa_prefill_identity: Optional[torch.Tensor] = None + # (identity table, (batch, rows), k2q CSR worklist) the NVFP4 prefill built for it: the worklist is + # layer-invariant for an identity selection, so the later layers of the step reuse it. + _msa_identity_csr: Optional[tuple] = None # This step's per-request host lengths, staged by _stage_host_lengths. _msa_qo_lens_cpu: Optional[torch.Tensor] = None _msa_kv_lens_cpu: Optional[torch.Tensor] = None @@ -650,6 +677,40 @@ def _ensure_prefill_n_valid_buffer(self, total_q: int, device: torch.device) -> self._msa_prefill_n_valid_buf = buf return buf + def _stage_msa_short_table(self, n_valid: torch.Tensor, params) -> None: + """Stage the identity block table for short-sequence steps (see msa_short_table).""" + self._msa_short_ok = False + if ( + not self.msa_policy.short_sequence_identity + or n_valid.numel() == 0 + or self.num_contexts > 0 + ): + return + topk = int(MSA_REQUIRED_TOPK) + if int(n_valid.max()) > topk: + return + _, num_kv_heads = params.sharded_head_counts(self.mapping) + total_q = int(n_valid.shape[0]) + max_tokens = max(int(self.max_num_sequences), self._msa_max_decode_tokens(), total_q) + if ( + self.msa_short_table is None + or self.msa_short_table.numel() < num_kv_heads * max_tokens * topk + ): + self.msa_short_table = self.get_empty( + self.cuda_graph_buffers, + (num_kv_heads * max_tokens * topk,), + cache_name="msa_short_table", + dtype=torch.int32, + capture_graph=self.is_cuda_graph, + ) + 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 + ) + self._msa_short_ok = True + def prepare(self) -> None: super().prepare() self._check_beam_width() @@ -847,6 +908,29 @@ def on_update_kv_lens(self) -> None: total_q = int(self.num_tokens) if batch <= 0 or total_q <= 0: return + page = int(self.kv_cache_manager.tokens_per_block) + n_valid_buf = ( + self._msa_prefill_n_valid_blocks + if self._msa_prefill_n_valid_blocks is not None + else self.msa_n_valid_blocks + ) + kv_update = self.msa_policy.kv_update + if kv_update is not None and kv_update( + self.kv_lens_cuda, + self.msa_kv_lens_staged, + self.seq_lens_cuda, + self.msa_q_batch_row, + self.msa_q_intra, + self.msa_block_table, + page, + batch, + total_q, + self.msa_seq_lens_cuda, + None if self.msa_cu_kv_lens is None else self.msa_cu_kv_lens[: batch + 1], + self.msa_out_cache_loc, + None if n_valid_buf is None else n_valid_buf[:total_q], + ): + return # kv_lens_cuda is the attended length (no num_extra_kv_tokens), the same # domain as the staged bound and msa_seq_lens_cuda. kv_true = torch.minimum(self.kv_lens_cuda[:batch], self.msa_kv_lens_staged[:batch]) @@ -907,6 +991,8 @@ def _build_step_plans(self) -> None: self._msa_prefill_gqa_plan = None self._msa_prefill_dense_plan = None self._msa_prefill_n_valid_blocks = None + self._msa_prefill_identity = None + self._msa_short_ok = False if not self._msa_fields_ready: return # Geometry is captured in __post_init__; skip when it is unavailable. @@ -937,7 +1023,9 @@ def _build_step_plans(self) -> None: qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size ) total_q = int(n_valid.shape[0]) - self.msa_n_valid_blocks[:total_q].copy_(n_valid.to(torch.int32), non_blocking=True) + # Packed but never skipped: on_update_kv_lens rewrites this buffer on the device. + step_inputs.copy_(self.msa_n_valid_blocks[:total_q], n_valid.to(torch.int32)) + self._stage_msa_short_table(n_valid, params) return fmha_sm100 = require_msa_module() @@ -968,13 +1056,28 @@ def plan_for(rows: Optional[Tuple[int, int]], **plan_kwargs) -> Optional[tuple]: **plan_kwargs, ) + n_valid_host = per_token_valid_blocks( + qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size + ) + total_q = int(n_valid_host.shape[0]) + short_prefill = ( + self.msa_policy.short_sequence_identity + and self._msa_decode_span is None + and total_q > 0 + and int(n_valid_host.max()) <= int(params.topk) + ) # Proxy plan: MQA (num_kv_heads=1) max-score pass over the index # branch; output_maxscore feeds the indexer's top-k block selection. - self._msa_prefill_proxy_plan = plan_for( - self._msa_proxy_plan_rows(), - num_qo_heads=num_index_heads, - num_kv_heads=1, - output_maxscore=True, + # A short pure-prefill step selects every valid block, so it needs none. + self._msa_prefill_proxy_plan = ( + None + if short_prefill + else plan_for( + self._msa_proxy_plan_rows(), + num_qo_heads=num_index_heads, + num_kv_heads=1, + output_maxscore=True, + ) ) attn_rows = self._msa_attn_plan_rows() # Sparse layers: kv_block_num=topk limits attention to top-k blocks. @@ -995,14 +1098,15 @@ def plan_for(rows: Optional[Tuple[int, int]], **plan_kwargs) -> Optional[tuple]: ) # Stage the valid-block count to the device once for the whole step # (see msa_prefill_n_valid_blocks). - n_valid_host = per_token_valid_blocks( - qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size - ) - total_q = int(n_valid_host.shape[0]) if total_q > 0: dev_buf = self._ensure_prefill_n_valid_buffer(total_q, _cache_device(self)) dev_buf[:total_q].copy_(n_valid_host.to(torch.int32), non_blocking=True) self._msa_prefill_n_valid_blocks = dev_buf[:total_q] + if short_prefill: + topk = int(params.topk) + ar = torch.arange(topk, dtype=torch.int32, device=dev_buf.device) + table = torch.where(ar[None, :] < dev_buf[:total_q, None], ar[None, :], -1) + self._msa_prefill_identity = table def _build_msa_fields(self) -> None: """Populate the MSA cache-write buffers for this step. @@ -1020,11 +1124,11 @@ def _build_msa_fields(self) -> None: kv_lens_cpu = self.msa_kv_lens_cpu qo_offset_cpu = self.msa_qo_offset_cpu if request_ids is None or qo_lens_cpu is None: - self.msa_out_cache_loc.fill_(-1) + _fill_unused_slots(self.msa_out_cache_loc, 0) return batch_size = int(qo_lens_cpu.shape[0]) if batch_size == 0: - self.msa_out_cache_loc.fill_(-1) + _fill_unused_slots(self.msa_out_cache_loc, 0) return kv_cache_manager = self.kv_cache_manager @@ -1041,6 +1145,8 @@ def _build_msa_fields(self) -> None: qo_lens_cpu=qo_lens_cpu, qo_offset_cpu=qo_offset_cpu, device=cache_device, + # Only out_cache_loc and the host block ids are read here; stage the slots straight from the host. + device_tables=not h2d_cache.enabled(), ) out_cache_loc = mapping.out_cache_loc # Only fmha_sm100 reads the flattened page table (the decode kernels @@ -1082,41 +1188,62 @@ def _build_msa_fields(self) -> None: f"smaller than the step's per-request page count ({block_table_cols})." ) - self.msa_out_cache_loc[:total_new_tokens].copy_(out_cache_loc, non_blocking=True) + step_inputs.copy_(self.msa_out_cache_loc[:total_new_tokens], out_cache_loc) # Captured producers also execute padded rows: the fused index producer # sits inside the captured region, so trimming it to a host-side count # 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. - if total_new_tokens < self.msa_out_cache_loc.shape[0]: - self.msa_out_cache_loc[total_new_tokens:].fill_(-1) + _fill_unused_slots(self.msa_out_cache_loc, total_new_tokens) if kv_indices is not None: - self.msa_kv_indices[: int(kv_indices.shape[0])].copy_(kv_indices, non_blocking=True) + step_inputs.copy_(self.msa_kv_indices[: int(kv_indices.shape[0])], kv_indices) # 2-D page table and per-request length for the decode kernels, # 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. - self.msa_block_table[:batch_size, :block_table_cols].copy_( - maybe_pin_memory(block_ids_cpu.to(torch.int32)), non_blocking=True + 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, ) - self.msa_seq_lens_cuda[:batch_size].copy_(kv_lens_cpu, non_blocking=True) + # 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) + else: + self.msa_seq_lens_cuda[:batch_size].copy_(kv_lens_cpu, non_blocking=True) if self.msa_cu_q_lens is not None: self._msa_live_batch = batch_size - self.msa_cu_q_lens[0].zero_() - self.msa_cu_kv_lens[0].zero_() qo_lens_cuda = self.msa_qo_lens_cuda[:batch_size] - qo_lens_cuda.copy_(qo_lens_cpu, non_blocking=True) - torch.cumsum( - qo_lens_cuda, - 0, - out=self.msa_cu_q_lens[1 : batch_size + 1], - ) - torch.cumsum( - self.msa_seq_lens_cuda[:batch_size], - 0, - out=self.msa_cu_kv_lens[1 : batch_size + 1], - ) + 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( + self.msa_cu_q_lens[: batch_size + 1], + _host_prefix_sum(qo_lens_cpu, self.msa_cu_q_lens.dtype), + pack=True, + ) + # Not through the skip cache: on_update_kv_lens rewrites this buffer on the device. + step_inputs.copy_( + self.msa_cu_kv_lens[: batch_size + 1], + _host_prefix_sum(kv_lens_cpu, self.msa_cu_kv_lens.dtype), + ) + else: + self.msa_cu_q_lens[0].zero_() + self.msa_cu_kv_lens[0].zero_() + qo_lens_cuda.copy_(qo_lens_cpu, non_blocking=True) + torch.cumsum( + qo_lens_cuda, + 0, + out=self.msa_cu_q_lens[1 : batch_size + 1], + ) + torch.cumsum( + self.msa_seq_lens_cuda[:batch_size], + 0, + out=self.msa_cu_kv_lens[1 : batch_size + 1], + ) self._msa_max_q_len = int(qo_lens_cpu.max().item()) self._msa_max_kv_len_all = int(kv_lens_cpu.max().item()) self._msa_total_k = int(kv_lens_cpu.to(torch.int64).sum().item()) @@ -1151,7 +1278,7 @@ 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: + if 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, @@ -1175,14 +1302,16 @@ def _build_msa_fields(self) -> None: torch.arange(total_new_tokens, dtype=torch.int64) - torch.repeat_interleave(starts, qo_long) ).to(torch.int32) - self.msa_q_batch_row[:total_new_tokens].copy_( - maybe_pin_memory(batch_row_cpu), non_blocking=True + copy_h2d_if_changed( + self.msa_q_batch_row[:total_new_tokens], maybe_pin_memory(batch_row_cpu), pack=True + ) + copy_h2d_if_changed( + self.msa_q_intra[:total_new_tokens], maybe_pin_memory(intra_cpu), pack=True ) - self.msa_q_intra[:total_new_tokens].copy_(maybe_pin_memory(intra_cpu), non_blocking=True) # The staged lens are the upper bound on_update_kv_lens clamps to: the # same attended lengths msa_seq_lens_cuda was just filled from, kept # apart so the clamp stays fixed however often the hook runs. - self.msa_kv_lens_staged[:batch_size].copy_(kv_lens_cpu, non_blocking=True) + step_inputs.copy_(self.msa_kv_lens_staged[:batch_size], kv_lens_cpu) self._msa_fields_ready = True def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor: @@ -1475,6 +1604,11 @@ def run_indexer( # proxy plan scores the whole batch instead. gen_first is the span's # first query token: the scorer takes [gen_first, num_tokens) over rows # [ctx_rows, row_last), the proxy the context prefix ahead of both. + identity = metadata._msa_prefill_identity + if span is None and identity is not None and int(identity.shape[0]) == num_tokens: + # Short pure-prefill step: every token selects all its valid blocks (staged once per step). + num_kv_heads = int(self.indexer.config.num_kv_heads) + return identity.unsqueeze(1).expand(num_tokens, num_kv_heads, int(identity.shape[1])) block_table = None seq_lens_cuda = None decode_query_len = None diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index d29def811bc5..8a262843975e 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -39,7 +39,9 @@ from tensorrt_llm.math_utils import ceil_div from tensorrt_llm.models.modeling_utils import QuantConfig +from ...pyexecutor import step_inputs from ...pyexecutor.config_utils import is_mla +from ...pyexecutor.h2d_cache import copy_h2d_if_changed from ...utils import (compute_swizzled_sf_shape, get_global_attrs, get_model_extra_attrs, helix_local_len_tensor) from .fmha.manager import FmhaManager @@ -997,8 +999,8 @@ def prepare(self) -> None: device='cpu', ) self.prompt_lens_cpu[:self.num_seqs].copy_(prompt_lens) - self.prompt_lens_cuda[:self.num_seqs].copy_( - self.prompt_lens_cpu[:self.num_seqs], non_blocking=True) + copy_h2d_if_changed(self.prompt_lens_cuda[:self.num_seqs], + self.prompt_lens_cpu[:self.num_seqs]) # number of tokens in the kv cache for each sequence in the batch cached_token_lens = torch.tensor( @@ -1051,9 +1053,8 @@ def prepare(self) -> None: # the sequence length including the cached tokens and the input tokens. self.kv_lens[:self.num_seqs].copy_( kv_lens + self.kv_cache_params.num_extra_kv_tokens) - self.kv_lens_cuda[:self.num_seqs].copy_(maybe_pin_memory( - kv_lens[:self.num_seqs]), - non_blocking=True) + step_inputs.copy_(self.kv_lens_cuda[:self.num_seqs], + maybe_pin_memory(kv_lens[:self.num_seqs])) # total kv lens for context requests and generation requests, without extra tokens self.host_total_kv_lens[0] = kv_lens[:self.num_contexts].sum().item() self.host_total_kv_lens[1] = kv_lens[self.num_contexts:self. @@ -1160,7 +1161,7 @@ def prepare_encoder_decoder_from_precomputed_lengths( f"({self.kv_cache_manager.max_seq_len}).") num_seqs = self.num_seqs - self.prompt_lens_cuda[:num_seqs].copy_(prompt_lens, non_blocking=True) + copy_h2d_if_changed(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/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index f01da1dc55ed..69cc41e928d4 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -9235,8 +9235,8 @@ def is_supported( @classmethod def _compile(cls, q_dtype: torch.dtype, num_heads: int, - max_decode_query_len: int, head_dim: int): - key = (q_dtype, num_heads, max_decode_query_len, head_dim) + max_decode_query_len: int, head_dim: int, split_k: int): + key = (q_dtype, num_heads, max_decode_query_len, head_dim, split_k) if key in cls.kernel_cache: return @@ -9283,7 +9283,7 @@ def _sym_stride(): cute_dtype, num_heads, max_decode_query_len, - cls.SPLIT_K, + split_k, head_dim, ) cls.kernel_cache[key] = cute.compile( @@ -9308,19 +9308,24 @@ def forward( seq_lens: torch.Tensor, score: torch.Tensor, max_decode_query_len: int, + split_k: int = SPLIT_K, ) -> None: """Score one decode step, compiling the kernel on first use. + ``split_k`` is the number of CTAs per request; each block is scored + by exactly one of them, so the scores do not depend on it. + The compile allocates, so it must not land inside a CUDA graph capture. It does not: CUDAGraphRunner.capture runs eager warmup forwards first, and those cover every geometry the graphs it then captures will replay. """ _, num_heads, head_dim = idx_q.shape - key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim) + key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim, + split_k) if key not in cls.kernel_cache: cls._compile(idx_q.dtype, num_heads, max_decode_query_len, - head_dim) + head_dim, split_k) cls.kernel_cache[key](idx_q, index_k_cache, block_table, score, seq_lens) @@ -9334,6 +9339,7 @@ def cute_dsl_minimax_m3_index_decode_score( seq_lens: torch.Tensor, score: torch.Tensor, max_decode_query_len: int, + split_k: int = 256, ) -> None: """Write per-block max index scores for one decode step, in place. @@ -9352,6 +9358,8 @@ def cute_dsl_minimax_m3_index_decode_score( passed without a copy. max_decode_query_len: compile-time bound on decode_query_len; num_index_heads * max_decode_query_len must not exceed 32. + split_k: CTAs per request (compile-time); each walks blocks + split, split + split_k, ... The scores do not depend on it. """ if not is_sm_100f(): raise ValueError( @@ -9376,6 +9384,7 @@ def cute_dsl_minimax_m3_index_decode_score( seq_lens, score, max_decode_query_len, + split_k, ) @torch.library.register_fake( @@ -9387,6 +9396,7 @@ def _( seq_lens: torch.Tensor, score: torch.Tensor, max_decode_query_len: int, + split_k: int = 256, ) -> None: 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 01b9f9cace7c..fb7c7092abce 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -827,6 +827,19 @@ def _choose_mxfp8_tactic(custom_op: str, runner: TunableRunner, return tactic +def mxfp8_quantize_autotuned( + input: torch.Tensor, + tune: bool = False) -> Tuple[torch.Tensor, torch.Tensor]: + """The quantize half of ``mxfp8_quantize_gemm_autotuned``: MXFP8 activation and scales, + on this token bucket's tuned quantize backend.""" + quantize_runner = MXFP8QuantizeRunner(input.dtype) + quantize_inputs = [input] + return quantize_runner(quantize_inputs, + tactic=_choose_mxfp8_tactic( + _MXFP8_QUANTIZE_AUTOTUNED_OP, quantize_runner, + quantize_inputs, tune)) + + def mxfp8_quantize_gemm_autotuned( input: torch.Tensor, weight: torch.Tensor, @@ -841,6 +854,24 @@ def mxfp8_quantize_gemm_autotuned( quantize_inputs, tactic=_choose_mxfp8_tactic(_MXFP8_QUANTIZE_AUTOTUNED_OP, quantize_runner, quantize_inputs, tune)) + return mxfp8_gemm_autotuned(act, act_scale, weight, weight_scale, + output_dtype, tune) + + +def mxfp8_gemm_autotuned( + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool = False, +) -> torch.Tensor: + """The GEMM half of mxfp8_quantize_gemm_autotuned for an activation already in MXFP8. + + Same runner, op key and tuning-cache entries as the decode-graph MXFP8 GEMM: + the tactic per token bucket comes from the AutoTuner (merged across ranks + when the runtime asks for it), never from a per-rank backend sweep. + """ gemm_runner = FlashInferMXFP8GemmRunner(output_dtype) gemm_inputs = [act, act_scale, weight, weight_scale] return gemm_runner(gemm_inputs, @@ -849,6 +880,43 @@ def mxfp8_quantize_gemm_autotuned( gemm_inputs, tune)) +def flashinfer_mxfp8_gemm_cute_dsl( + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + sweep_configs: bool, +) -> torch.Tensor: + """FlashInfer's CuTeDSL MXFP8 GEMM for an activation its producer already quantized. + + Unlike ``mxfp8_quantize_gemm_autotuned``, the backend is fixed rather than + chosen per token bucket: the AutoTuner profiles with the weight L2-resident, + where CUTLASS can edge out CuTeDSL, while a decode step streams it cold. + + ``sweep_configs`` runs FlashInfer's own CuTeDSL config sweep: it profiles + the configs for this shape while ``flashinfer_mxfp8_autotune`` is active + (the warmup tuning pass) and replays the cached best one afterwards; + without it the heuristic config runs. + """ + if sweep_configs: + from ..modules.linear import flashinfer_mxfp8_autotune_active + context = _flashinfer_autotune( + tune_mode=flashinfer_mxfp8_autotune_active()) + else: + context = _flashinfer_autotune(tune_mode=False, skip_ops="mxfp8_gemm") + with context: + return _flashinfer_mm_mxfp8( + act, + weight.t(), + act_scale, + weight_scale, + out_dtype=output_dtype, + use_8x4_sf_layout=False, + backend="cute-dsl", + ) + + class FP4GemmRunner(TunableRunner): runner_dict = dict() tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec( 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..b1c9d4de0e08 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/m3_mx_gemv_kernel.py @@ -0,0 +1,1236 @@ +# 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) + + +@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/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..abcbadff4721 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/op.py @@ -0,0 +1,474 @@ +# 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], +) -> 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 + 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/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 26d09566eba2..5a66d30efdcf 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -3366,6 +3366,16 @@ def _mxfp8_cutlass_op_available() -> bool: "flashinfer_mxfp8_decode_graph_capture_active", default=False) +def flashinfer_mxfp8_autotune_active() -> bool: + """Whether the current forward runs inside ``flashinfer_mxfp8_autotune``.""" + return _FLASHINFER_MXFP8_AUTOTUNE_ACTIVE.get() + + +def flashinfer_mxfp8_decode_graph_capture_active() -> bool: + """Whether the current forward is a decode CUDA-graph warmup or capture.""" + return _FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE.get() + + @contextmanager def flashinfer_mxfp8_autotune() -> Iterator[None]: """Tune FlashInfer MXFP8 tactics while enabling auto-dispatched calls.""" diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 7aedaf8a81e6..4c7d809309fb 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -51,6 +51,7 @@ from ..speculative import (draft_prompt_lookahead, get_num_extra_kv_tokens, get_num_spec_layers, get_spec_decoder, should_use_separate_draft_kv_cache) +from ..speculative.spec_sampler_base import SpecSampler from ..utils import is_gdn_replay_enabled from . import config_utils from .config_utils import (MambaKVCacheParams, _is_sliding_attention_layer, @@ -75,6 +76,7 @@ from .resource_manager import (KVCacheCompressionManager, KVCacheManager, PeftCacheManager, ResourceManager, ResourceManagerType) +from .runtime_hooks import runtime_hooks_of from .sampler import EarlyStopSampler, EarlyStopWithMMResult, TorchSampler from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler, KVCacheV2Scheduler, MultimodalEagerEncoderScheduler, @@ -4417,7 +4419,11 @@ def instantiate_sampler( ) if engine.spec_config is not None and engine.spec_config.spec_dec_mode.has_spec_decoder( ): - return get_spec_decoder(sampler_args, engine.spec_config) + 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) + return sampler if mm_encoder_only: # NOTE: handle model outputs specially for mm encoder executor/engine. diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 19ed990f33a1..2285ee60c0a6 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -27,6 +27,7 @@ from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..speculative.utils import get_draft_kv_cache_manager from ..utils import make_weak_ref, piecewise_cuda_graph +from . import step_inputs from .llm_request import LlmRequest, get_draft_token_length from .resource_manager import (BaseResourceManager, ResourceManager, ResourceManagerType) @@ -149,6 +150,15 @@ class CUDAGraphRunnerConfig: kv_cache_manager_key: Any dynamic_draft_len_mapping: Optional[Dict[int, int]] = None sparse_attention_config: Optional[BaseSparseAttentionConfig] = None + # Total sequence length at or below which a generation batch replays the + # short-sequence graph family; None captures one family. Derived from the + # sparse attention config when it asks for the split, else from the model's + # runtime hooks. + short_long_seq_len_threshold: Optional[int] = None + # Give the graph metadata a buffer pool of its own instead of the + # process-wide one, so no other engine writes its metadata buffers (an + # engine that skips unchanged uploads into them relies on that). + private_metadata_buffers: bool = False enable_encoder_decoder_mixed_cuda_graph: bool = False enable_in_graph_sampling: bool = False @@ -162,6 +172,8 @@ class CUDAGraphRunner: multiple graphs, keyed by batch shape and execution-path specializations. """ WARMUP_STEPS = 1 + # CUDAGraphRunnerConfig.short_long_seq_len_threshold, when the config set one. + _configured_seq_len_threshold: Optional[int] = None def __init__(self, config: CUDAGraphRunnerConfig): self.config = config @@ -174,6 +186,7 @@ def __init__(self, config: CUDAGraphRunnerConfig): self.max_beam_width = config.max_beam_width self.spec_config = config.spec_config self.sparse_config = config.sparse_attention_config + self._configured_seq_len_threshold = config.short_long_seq_len_threshold self.is_encoder_decoder = config.is_encoder_decoder self.enable_encoder_decoder_mixed_cuda_graph = ( config.enable_encoder_decoder_mixed_cuda_graph) @@ -203,7 +216,9 @@ def __init__(self, config: CUDAGraphRunnerConfig): self.shared_static_tensors: Dict[str, torch.Tensor] = {} if self.enabled: self._create_shared_static_tensors() - self.cuda_graph_meta_buffers = get_memory_buffers() + self.cuda_graph_meta_buffers = (Buffers() + if config.private_metadata_buffers else + get_memory_buffers()) # On-the-fly capture is disabled by default to prevent workspace # tensor reallocation from invalidating addresses baked into existing @@ -301,6 +316,20 @@ def _gather_adp_graph_batch_info(self, batch: ScheduledRequests, tiers = gathered[:, 2].tolist() return list(zip(flags, sizes, tiers)) + @property + def short_long_seq_len_threshold(self) -> Optional[int]: + """The total sequence length splitting the short and long graph families, or None for one family. + + The runner config's value when it set one, else the sparse attention + config's when that asks for the split. + """ + if self._configured_seq_len_threshold is not None: + return self._configured_seq_len_threshold + if (isinstance(self.sparse_config, SeqLenAwareSparseAttentionConfig) + and self.sparse_config.needs_separate_short_long_cuda_graphs()): + return self.sparse_config.seq_len_threshold + return None + def _get_seq_len_mode( self, batch: ScheduledRequests, @@ -313,8 +342,7 @@ def _get_seq_len_mode( that the model engine temporarily placed in the generation list. It is empty for the existing generation-only path. """ - if (isinstance(self.sparse_config, SeqLenAwareSparseAttentionConfig) - and self.sparse_config.needs_separate_short_long_cuda_graphs()): + if self.short_long_seq_len_threshold is not None: # Some sparse attention algorithms need to use different forward paths for short and long sequences. # For example, the DSA can skip the MQA and Top-K in the indexer for short sequences to reduce the # computational overhead. To support this feature, we need to capture separate CUDA graphs for short @@ -356,7 +384,7 @@ def _get_seq_len_mode( from ..speculative import get_num_extra_kv_tokens num_extra_kv_tokens = get_num_extra_kv_tokens(self.spec_config) max_seq_len = max(total_seq_lens) - if max_seq_len <= self.sparse_config.seq_len_threshold - num_extra_kv_tokens: + if max_seq_len <= self.short_long_seq_len_threshold - num_extra_kv_tokens: short_seq_len_mode = True else: short_seq_len_mode = False @@ -816,7 +844,7 @@ def replay( f"replay() got {seqlen} tokens for key {key}, but the graph " f"was captured for {expected_num_tokens} tokens. A shorter " "input_ids leaves the tail of the static input buffer stale.") - static_tensors["input_ids"][:seqlen].copy_(input_ids) + step_inputs.copy_(static_tensors["input_ids"][:seqlen], input_ids) position_ids = current_inputs["position_ids"] if self.config.use_mrope: @@ -863,7 +891,8 @@ def replay( f"for key {key}, but expected {expected_position_ids_shape}. " "torch.Tensor.copy_() silently broadcasts mismatched shapes, " "which would corrupt the static input buffer.") - static_tensors["position_ids"][:, :seqlen].copy_(position_ids) + step_inputs.copy_(static_tensors["position_ids"][:, :seqlen], + position_ids) num_encoder_tokens = key.num_encoder_tokens if num_encoder_tokens: @@ -887,6 +916,9 @@ def replay( encoder_hidden_states) 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() if self._pre_replay_hooks: hook_tensors = static_tensors if position_id_offsets is not None: diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 3b8d7fa2da91..f533101257f4 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -34,6 +34,7 @@ ) from tensorrt_llm._torch.peft.lora.config import LoraConfig from tensorrt_llm._torch.peft.lora.manager import LoraModelConfig +from tensorrt_llm._torch.pyexecutor import h2d_cache, step_inputs from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner import BreakableCUDAGraphRunner from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( CUDAGraphRunner, @@ -42,6 +43,7 @@ get_mrope_dummy_seq_slot, ) from tensorrt_llm._torch.pyexecutor.guided_decoder import CapturableGuidedDecoder +from tensorrt_llm._torch.pyexecutor.h2d_cache import H2DSkipCache, copy_h2d_if_changed from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( BaseMambaCacheManager, @@ -54,11 +56,23 @@ ResourceManager, ResourceManagerType, ) +from tensorrt_llm._torch.pyexecutor.runtime_hooks import ( + DEFAULT_RUNTIME_HOOKS, + MsaRuntimePolicy, + RuntimeHooks, + runtime_hooks_of, +) from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import warmup_sampling_module +from tensorrt_llm._torch.pyexecutor.sampler.sampler import ( + GREEDY_HIDDEN_STATES_KEY, + GREEDY_TOKENS_KEY, + greedy_tokens_usable, +) from tensorrt_llm._torch.pyexecutor.sampler.sampler_common import SampleType from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import RemoteTailPhase +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack from tensorrt_llm._torch.pyexecutor.trace_log_utils import log_mem_snapshot from tensorrt_llm._torch.pyexecutor.warmup_timer import _WarmupTimer from tensorrt_llm._torch.pyexecutor.workspace import EagerWorkspaceReclaimer @@ -91,6 +105,7 @@ from tensorrt_llm.llmapi.llm_args import ( BaseSparseAttentionConfig, DecodingBaseConfig, + MiniMaxM3SparseAttentionConfig, PrefillCudaGraphBackend, SeqLenAwareSparseAttentionConfig, ) @@ -296,6 +311,15 @@ def __getattr__(self, name: str) -> Any: class DecoderRunner(ScheduledModelRunner): """Run scheduled decoder generation with its own buffers, graphs and metadata.""" + # Step-loop behaviour the model opted into (RuntimeHooks); all off unless it declares otherwise. + runtime_hooks: RuntimeHooks = DEFAULT_RUNTIME_HOOKS + # This engine's skip cache and step-input pack (None: the hook is off), active only during its own forward. + _h2d_cache: H2DSkipCache | None = None + _step_pack: StepInputPack | None = None + # Whether a batch may keep the greedy tokens a forward sampled in place of the logits. False when the executor + # runs a guided decoder, which reads the logits of every batch. + keep_greedy_only_outputs = True + # Optional execution paths that a subclass can turn off. _context_warmups_supported = True _context_graph_promotion_supported = True @@ -347,6 +371,13 @@ def __init__( self._mm_item_scheduler = mm_item_scheduler self._config = config + self.runtime_hooks = runtime_hooks_of(model) + self._h2d_cache = H2DSkipCache() if self.runtime_hooks.skip_unchanged_h2d else None + self._step_pack = ( + StepInputPack(torch.device("cuda", torch.cuda.current_device())) + if self.runtime_hooks.pack_step_inputs + else None + ) self.guided_decoder: CapturableGuidedDecoder | None = None self.lora_model_config: LoraModelConfig | None = None @@ -391,7 +422,15 @@ def forward( resource_manager: ResourceManager, is_dummy: bool = False, ) -> Any: - return self._forward_decoder(inputs, resource_manager, is_dummy=is_dummy) + with h2d_cache.active(self._h2d_cache), step_inputs.active(self._step_pack): + try: + return self._forward_decoder(inputs, resource_manager, is_dummy=is_dummy) + except BaseException: + # A failed step issues none of its packed copies, so forget the shadows recorded for them too. + step_inputs.discard() + if self._h2d_cache is not None: + self._h2d_cache.clear() + raise def warmup(self, resource_manager: ResourceManager) -> None: kv_cache_manager = resource_manager.get_resource_manager(self.kv_cache_manager_key) @@ -472,6 +511,14 @@ def _allocate_decoder_buffers(self) -> None: self.position_ids_cuda = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, device="cuda" ) + if self._config.is_spec_decode and self.runtime_hooks.spec_step_fold is not None: + self.runtime_hooks.spec_step_fold.bind_engine_buffers( + self.input_ids_cuda, + self.draft_tokens_cuda, + self.previous_pos_id_offsets_cuda, + self.previous_kv_lens_offsets_cuda, + self._config.max_batch_size, + ) # Host counter of the steady-state generation prepare cache. self._steady_gen_positions_pinned = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, pin_memory=prefer_pinned() @@ -550,9 +597,33 @@ def _cuda_graph_runner_config(self) -> CUDAGraphRunnerConfig: dist=self.dist, kv_cache_manager_key=self.kv_cache_manager_key, sparse_attention_config=self.sparse_attention_config, + short_long_seq_len_threshold=self._short_long_seq_len_threshold(), + private_metadata_buffers=self.runtime_hooks.skips_unchanged_metadata_uploads, enable_in_graph_sampling=self._config.enable_in_graph_sampling, ) + def _short_long_seq_len_threshold(self) -> int | None: + """The total sequence length that splits the short and long decode graph families, or None for one family. + + 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. + """ + sparse_config = self.sparse_attention_config + if ( + isinstance(sparse_config, SeqLenAwareSparseAttentionConfig) + and sparse_config.needs_separate_short_long_cuda_graphs() + ): + return sparse_config.seq_len_threshold + if self.runtime_hooks.msa.short_sequence_identity: + if not isinstance(sparse_config, MiniMaxM3SparseAttentionConfig): + raise ValueError( + "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 + return None + def _initialize_breakable_cuda_graph_runner(self) -> BreakableCUDAGraphRunner | None: if ( self.cuda_graph_runner is None @@ -1643,6 +1714,11 @@ def run_autotuner_pass(autotune_context: Any, synchronize_trtllm_cache: bool) -> ran_flashinfer_forward = run_autotuner_pass( flashinfer_mxfp8_autotune(), synchronize_trtllm_cache=False ) + if self.runtime_hooks.merge_tuned_tactics_across_ranks: + # Rank-local profiling noise must not leave tensor-parallel peers + # serving different tactics for one shape: keep the fastest everywhere. + # Collective; before the native MXFP8 cache below registers the tactics. + AutoTuner.get().post_tune_merge_tactics() if enable_flashinfer_mxfp8_autotuner: if ran_flashinfer_forward: @@ -2037,6 +2113,10 @@ def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): ) with flashinfer_autotune_context, flashinfer_mxfp8_decode_graph_capture(): self._capture_generation_cuda_graphs(resource_manager) + if is_warmup_only and self.runtime_hooks.merge_tuned_tactics_across_ranks: + # Tactics tuned while warming the decode graphs agree across + # ranks before the capture pass reads them. + AutoTuner.get().post_tune_merge_tactics() self._capture_additional_cuda_graphs(resource_manager) # Piecewise graphs have separate capture machinery and do not use the # whole-model attention workspace. Capture them only on the second pass. @@ -2081,11 +2161,8 @@ def _capture_generation_cuda_graphs(self, resource_manager: ResourceManager): if self.mapping is not None and self.mapping.has_cp_helix(): effective_max_seq_len = self._config.max_seq_len // self.mapping.cp_size - sparse_config = self.sparse_attention_config - if ( - isinstance(sparse_config, SeqLenAwareSparseAttentionConfig) - and sparse_config.needs_separate_short_long_cuda_graphs() - ): + short_long_seq_len_threshold = self._short_long_seq_len_threshold() + if short_long_seq_len_threshold is not None: # For short sequences, subtract the maximum runtime tokens consumed # by a generation step so all current-step tokens stay within the # sequence length threshold. PARD uses 2K tokens here, not K+1. @@ -2093,7 +2170,7 @@ def _capture_generation_cuda_graphs(self, resource_manager: ResourceManager): self._config.max_draft_len ) # For long sequences, use the default maximum sequence length. - max_seq_len = sparse_config.seq_len_threshold - max_runtime_tokens_per_gen_step + max_seq_len = short_long_seq_len_threshold - max_runtime_tokens_per_gen_step if max_seq_len < effective_max_seq_len: max_seq_len_list = [effective_max_seq_len, max_seq_len] else: @@ -2776,15 +2853,43 @@ def _set_up_attn_metadata( config, "chunk_size", self.attn_metadata.mamba_chunk_size ) self.attn_metadata.mamba_metadata_cls = resolve_mamba_metadata_cls(self.model) + self._arm_msa_policy(kv_cache_manager) return self.attn_metadata + def _arm_msa_policy(self, kv_cache_manager: Any) -> None: + """Hand the model's MiniMax-M3 MSA preparation choices to the metadata and the KV cache.""" + msa = self.runtime_hooks.msa + if msa == MsaRuntimePolicy(): + return + # Imported here: only a model that armed an MSA policy needs the sparse backend. + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager import ( + MiniMaxM3KVCacheManagerV2, + ) + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, + ) + + if not isinstance( + self.attn_metadata, MiniMaxM3MsaSparseAttentionMetadata + ) or not isinstance(kv_cache_manager, MiniMaxM3KVCacheManagerV2): + raise ValueError( + "RuntimeHooks.msa requires the MiniMax-M3 MSA sparse attention backend and KV cache, " + f"got {type(self.attn_metadata).__name__} and {type(kv_cache_manager).__name__}" + ) + self.attn_metadata.msa_policy = msa + kv_cache_manager.configure_decode_shortcuts( + skip_unchanged_block_offsets=msa.skip_unchanged_block_offsets, + reuse_pool_views=msa.reuse_pool_views, + ) + def _set_up_spec_metadata(self, spec_resource_manager: BaseResourceManager | None): if self.spec_metadata is not None: return self.spec_metadata self.spec_metadata = create_spec_metadata( self._config, self.model.config, spec_resource_manager ) + self.spec_metadata.skip_greedy_rng_upload = self.runtime_hooks.skip_greedy_rng_upload return self.spec_metadata def _release_decoder_graphs(self) -> None: @@ -4002,7 +4107,7 @@ def previous_seq_slots_device(): previous_batch_indices, dtype=torch.int, pin_memory=prefer_pinned() ) previous_slots = self.previous_batch_indices_cuda[:previous_batch_len] - previous_slots.copy_(previous_batch_indices_host, non_blocking=True) + copy_h2d_if_changed(previous_slots, previous_batch_indices_host) return previous_slots if scheduled_requests.is_attention_dp_phase_idle: @@ -4027,17 +4132,68 @@ def previous_seq_slots_device(): num_accepted_draft_tokens = torch.tensor( num_accepted_draft_tokens, dtype=torch.int, pin_memory=prefer_pinned() ) - self.num_accepted_draft_tokens_cuda[: len(num_accepted_draft_tokens)].copy_( - num_accepted_draft_tokens, non_blocking=True + copy_h2d_if_changed( + self.num_accepted_draft_tokens_cuda[: len(num_accepted_draft_tokens)], + num_accepted_draft_tokens, + ) + step_fold = self.runtime_hooks.spec_step_fold if self._config.is_spec_decode else None + fold_rows = None + fold_staged = False + if step_fold is not None: + # A decode-only step's generation rows; the model's forward may stage the next + # step's inputs for exactly these rows (see SpecStepFold). + fold_rows = ( + [request.py_seq_slot for request in extend_requests] + if num_ctx_requests == 0 + else None ) + if fold_rows is not None and any(slot is None for slot in fold_rows): + fold_rows = None + fold_staged = step_fold.begin_step(fold_rows, len(extend_requests)) if next_draft_tokens_device is not None: - # Initialize these two values to zeros - self.previous_pos_id_offsets_cuda *= 0 - self.previous_kv_lens_offsets_cuda *= 0 runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step(runtime_draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 + # The previous forward already staged this step's inputs when it ran for exactly + # these rows and this step adds nothing of its own. + inputs_staged = ( + fold_staged + and len(extend_dummy_requests) == 0 + and num_tokens == 0 + and num_draft_tokens == 0 + and previous_batch_len == len(extend_requests) + and previous_batch_indices == fold_rows + ) + if not inputs_staged and step_fold is not None and previous_batch_len > 0: + previous_slots = previous_seq_slots_device() + previous_pos_indices_host = torch.tensor( + previous_pos_indices, dtype=torch.int, pin_memory=prefer_pinned() + ) + copy_h2d_if_changed( + self.previous_pos_indices_cuda[ + : previous_batch_len * runtime_tokens_per_gen_step + ], + previous_pos_indices_host, + ) + num_extend_reqeust_wo_dummy = len(extend_requests) - len(extend_dummy_requests) + inputs_staged = step_fold.gather( + new_tokens_device, + next_draft_tokens_device, + new_tokens_lens_device, + previous_slots, + self.previous_pos_indices_cuda, + previous_batch_len, + runtime_tokens_per_gen_step, + runtime_draft_token_buffer_width, + num_tokens, + num_draft_tokens, + num_extend_reqeust_wo_dummy - previous_batch_len, + ) + if not inputs_staged: + # Initialize these two values to zeros + self.previous_pos_id_offsets_cuda *= 0 + self.previous_kv_lens_offsets_cuda *= 0 - if previous_batch_len > 0: + if previous_batch_len > 0 and not inputs_staged: previous_slots = previous_seq_slots_device() # previous input ids previous_batch_tokens = previous_batch_len * runtime_tokens_per_gen_step @@ -4064,8 +4220,9 @@ def previous_seq_slots_device(): previous_pos_indices_host = torch.tensor( previous_pos_indices, dtype=torch.int, pin_memory=prefer_pinned() ) - self.previous_pos_indices_cuda[0:previous_batch_tokens].copy_( - previous_pos_indices_host, non_blocking=True + copy_h2d_if_changed( + self.previous_pos_indices_cuda[0:previous_batch_tokens], + previous_pos_indices_host, ) # The order of requests in a batch: [context requests, generation requests] @@ -4164,13 +4321,13 @@ def previous_seq_slots_device(): ) final_position_ids = self.mrope_position_ids_cuda[:, :, :total_num_tokens] else: - self.position_ids_cuda[:total_num_tokens].copy_(host_position_ids, non_blocking=True) + step_inputs.copy_(self.position_ids_cuda[:total_num_tokens], host_position_ids) final_position_ids = self.position_ids_cuda[:total_num_tokens].unsqueeze(0) if enable_spec_decode: - self.gather_ids_cuda[: len(gather_ids)].copy_( + copy_h2d_if_changed( + self.gather_ids_cuda[: len(gather_ids)], torch.tensor(gather_ids, dtype=torch.int, pin_memory=prefer_pinned()), - non_blocking=True, ) if self.mapping.has_cp_helix(): @@ -4790,6 +4947,16 @@ def _forward_decoder( execution_requests.generation_requests, ) + # Pack the step's small input copies into one upload and one scatter at the + # graph replay. Only a replay of an already captured graph, which always + # reaches the flush, may open a pack. + if ( + self.runtime_hooks.pack_step_inputs + and can_run_graph + and not self.cuda_graph_runner.is_warmup_only + and not self.cuda_graph_runner.needs_capture(key) + ): + step_inputs.begin() with self._cuda_graph_metadata_prepare_scope( attn_metadata, key, @@ -4921,9 +5088,14 @@ def capture_postprocess_fn(inputs: dict[str, Any]): finally: restore_attn_metadata_after_draft_replay(attn_metadata, saved_draft) + # No-op unless a pack was opened above and the replay did not issue it. + step_inputs.flush() + if self.forward_pass_callable is not None: self.forward_pass_callable() + if isinstance(outputs, dict) and outputs.get(GREEDY_TOKENS_KEY) is not None: + outputs = self._logits_for_greedy_tokens(semantic_requests, outputs) self._execute_logit_post_processors(semantic_requests, outputs) if not isinstance(outputs, dict): @@ -5044,6 +5216,39 @@ def _apply_logits_processors( # back would be a self-assignment, which torch rejects for the # non-contiguous slices a TP-padded vocab produces. + def _logits_for_greedy_tokens( + self, scheduled_requests: ScheduledRequests, outputs: dict + ) -> dict: + """Replace greedy tokens sampled by the forward with logits when a request needs logits. + + A forward that samples greedily in place of the LM head returns the head's input + alongside the tokens. When any generation request needs more than its greedy token, + the logits come from that input eagerly, followed by the in-graph sampling and the + logits post-processing the forward skipped. Returns the outputs to use: ``outputs`` + itself, or a new dict (a replayed graph's outputs dict is reused by later replays). + """ + if ( + self.keep_greedy_only_outputs + and self.max_beam_width == 1 + and scheduled_requests.num_context_requests == 0 + # An idle attention-DP rank runs a dummy row for an empty batch, whose outputs are restored from logits. + and len(scheduled_requests.generation_requests) > 0 + and all( + greedy_tokens_usable(request) for request in scheduled_requests.generation_requests + ) + ): + return outputs + outputs = dict(outputs) + outputs.pop(GREEDY_TOKENS_KEY) + hidden_states = outputs.pop(GREEDY_HIDDEN_STATES_KEY) + # One row per generation request: every row is a last token. + outputs["logits"] = self.model.logits_processor.forward( + hidden_states, self.model.lm_head, None, True + ) + if self.sample_in_graph_callable is not None and self.mapping.is_last_pp_rank(): + self.sample_in_graph_callable(outputs) + return outputs + def _execute_logit_post_processors(self, scheduled_requests: ScheduledRequests, outputs: dict): """Apply logit post processors (in-place modify outputs Tensors) if any.""" @@ -5055,6 +5260,10 @@ def _execute_logit_post_processors(self, scheduled_requests: ScheduledRequests, return logits_tensor = outputs["logits"] + if logits_tensor is None: + # Greedy tokens sampled by the forward; _logits_for_greedy_tokens supplied + # logits instead whenever a request needed them. + return logits_row_offset = 0 request_groups = ( diff --git a/tensorrt_llm/_torch/pyexecutor/h2d_cache.py b/tensorrt_llm/_torch/pyexecutor/h2d_cache.py new file mode 100644 index 000000000000..1695ac2e1189 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/h2d_cache.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Skip host-to-device metadata copies whose values the destination already holds. + +Steady decode re-uploads many small per-step metadata tensors whose values do not change (slot indices, query +lengths, prompt lengths, page tables between page boundaries). Each upload is a separate copy on the compute stream, +serialized between the decode graphs. An ``H2DSkipCache`` keeps one host shadow per destination buffer (the last +values written through it, and the view they were written to) and skips the copy when the same view receives the same +values again. + +Each engine that armed ``RuntimeHooks.skip_unchanged_h2d`` owns one cache and makes it the active one for its forward +(``active``); ``copy_h2d_if_changed`` acts on the calling context's active cache and, with none, is +``dst.copy_(src, non_blocking=True)``. The attention metadata that calls it has no engine to ask, which is why the +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. +""" + +from __future__ import annotations + +import contextlib +from contextvars import ContextVar +from typing import Iterator, Optional + +import torch + +from . import step_inputs + + +class H2DSkipCache: + """One engine's host shadows of its skip-routed destinations (see the module docstring).""" + + def __init__(self) -> None: + # 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. + self._shadow: dict[int, tuple[object, tuple, torch.Tensor]] = {} + # Device storage address -> (the storage, kept alive as above; the view key; (first row, value) of the view's + # last tail fill). + self._tails: dict[int, tuple[object, tuple, tuple[int, int]]] = {} + + def clear(self) -> None: + """Forget every shadow, so the next copy to each destination is issued (e.g. after a discarded step).""" + self._shadow.clear() + self._tails.clear() + + def copy(self, dst: torch.Tensor, src: torch.Tensor, pack: bool) -> bool: + storage = dst.untyped_storage() + addr = storage.data_ptr() + if src.is_cuda or not dst.is_cuda or torch.cuda.is_current_stream_capturing(): + self._shadow.pop(addr, None) + dst.copy_(src, non_blocking=True) + return True + key = (dst.storage_offset(), tuple(dst.shape), tuple(dst.stride()), dst.dtype) + entry = self._shadow.get(addr) + if ( + entry is not None + and entry[1] == key + and entry[2].dtype == src.dtype + and entry[2].shape == src.shape + and torch.equal(entry[2], src) + ): + return False + if pack: + step_inputs.copy_(dst, src) + else: + dst.copy_(src, non_blocking=True) + self._shadow[addr] = (storage, key, src.clone()) + return True + + def fill_tail(self, buf: torch.Tensor, start: int, value: int) -> None: + storage = buf.untyped_storage() + addr = storage.data_ptr() + if start >= buf.shape[0]: + # The live rows cover the whole view: the previous tail fill is overwritten. + self._tails.pop(addr, None) + return + if torch.cuda.is_current_stream_capturing(): + self._tails.pop(addr, None) + buf[start:].fill_(value) + return + key = (buf.storage_offset(), tuple(buf.shape), tuple(buf.stride()), buf.dtype) + entry = self._tails.get(addr) + if entry is not None and entry[1] == key and entry[2] == (start, value): + return + buf[start:].fill_(value) + self._tails[addr] = (storage, key, (start, value)) + + +_active: ContextVar[Optional[H2DSkipCache]] = ContextVar("trtllm_h2d_skip_cache", default=None) + + +@contextlib.contextmanager +def active(cache: Optional[H2DSkipCache]) -> Iterator[None]: + """Make ``cache`` the calling context's cache (None: no skipping) for the duration.""" + token = _active.set(cache) + try: + yield + finally: + _active.reset(token) + + +def enabled() -> bool: + """Whether the calling context skips unchanged copies.""" + return _active.get() is not None + + +def copy_h2d_if_changed(dst: torch.Tensor, src: torch.Tensor, pack: bool = False) -> bool: + """``dst.copy_(src, non_blocking=True)`` unless the active cache knows ``dst`` holds ``src``; True if copied. + + ``pack``: the copy may join an open ``step_inputs`` pack, i.e. land only at the pack's flush (the graph replay). + Pass it only for destinations no device work reads before that replay. + """ + cache = _active.get() + if cache is None: + dst.copy_(src, non_blocking=True) + return True + return cache.copy(dst, src, pack) + + +def fill_tail_if_changed(buf: torch.Tensor, start: int, value: int) -> None: + """``buf[start:].fill_(value)`` unless the active cache's last tail fill of ``buf`` was the same. + + Call it after every write of the live rows ``buf[:start]``, including one that covers the whole of ``buf`` + (``start == len(buf)``, which fills nothing): a live-row write past the recorded start overwrites part of the + tail. The skip is sound only while nothing else writes ``buf`` past ``start``. + """ + cache = _active.get() + if cache is None: + if start < buf.shape[0]: + buf[start:].fill_(value) + return + cache.fill_tail(buf, start, value) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 221f9145b17c..d87da7f1e6be 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -854,6 +854,11 @@ def _init_cuda_graph_lora_manager(self, lora_config: LoraConfig): if isinstance(self._runner, DecoderRunner): self._runner.init_cuda_graph_lora_manager(lora_config) + def drop_greedy_only_outputs(self) -> None: + """Have every batch return logits, for an executor-side guided decoder that reads them.""" + if isinstance(self._runner, DecoderRunner): + self._runner.keep_greedy_only_outputs = False + def set_guided_decoder(self, guided_decoder: CapturableGuidedDecoder) -> bool: if hasattr(self.model, "set_guided_decoder"): diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index a88984ea9ee5..3d6be068f311 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -114,6 +114,7 @@ get_from_waiting_queue, merge_requests) from .resource_manager import (NoFreeSlotsError, ResourceManager, ResourceManagerType, request_context) +from .runtime_hooks import runtime_hooks_of from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState, SampleStateTensors) from .scheduler import (RequestScheduler, ScheduledRequests, @@ -405,6 +406,13 @@ class PendingEncoderStep: class PyExecutor: + # True while the overlap loop runs with the execution stream current + # (_overlap_loop_stream_scope). + _loop_on_execution_stream = False + # Whether the model opted into running the overlap loop on the execution + # stream (RuntimeHooks.loop_on_execution_stream); read once at construction. + _execution_stream_loop_armed = False + # Minimum number of async micro batches for async PP execution. # This is a trade-off between memory usage and performance. # If the number of micro batches is too small, the executor will spend too much time in synchronization. @@ -516,6 +524,8 @@ def __init__( # related modules self.resource_manager = resource_manager self.model_engine = model_engine + self._execution_stream_loop_armed = runtime_hooks_of( + model_engine.model).loop_on_execution_stream self._enable_adp_dummy_fixes = getattr(model_engine, "_enable_adp_dummy_fixes", False) self._enable_scheduler_aware_adp_dummy = getattr( @@ -5693,11 +5703,42 @@ def _wait_for_model_engine_input_copy(self): if wait_for_input_copy is not None: wait_for_input_copy() + @contextmanager + def _overlap_loop_stream_scope(self): + """Run the overlap loop with the execution stream current, when the model opted in. + + _forward_step makes the execution stream wait for the current stream + before every forward and the current stream wait for the execution + stream after it. On the thread's default stream, which in steady decode + carries only these waits and timing events, step N+1's inputs wait for + decode graph N through two cross-stream event hops. With the execution + stream current the loop's stream work is queued in the same order on one + stream and _forward_step skips both waits. + + Not taken with a two-model drafter, a guided decoder, a KV connector, + an encoder-decoder model or multimodal encoder item scheduling, whose + side streams synchronize with the current stream. + """ + if not (self._execution_stream_loop_armed and self.drafter is None and + self.draft_model_engine is None and self.guided_decoder is None + and self.kv_connector_manager is None + and not self.is_encoder_decoder + and not self._mm_encoder_item_scheduling_enabled): + yield + return + with torch.cuda.stream(self.execution_stream): + self._loop_on_execution_stream = True + try: + yield + finally: + self._loop_on_execution_stream = False + def _executor_loop_overlap(self): torch.cuda.set_device(self.device_id) # ensure the context is created, otherwise, some MPI calls will fail. CUASSERT(cudart.cudaSetDevice(self.device_id)) - with self._profiler() as profile_step, self.hang_detector: + with self._profiler() as profile_step, self.hang_detector, \ + self._overlap_loop_stream_scope(): iter_start_time = time.time() iter_stats = None target_inputs = None @@ -8402,8 +8443,10 @@ def forward(scheduled_requests, resource_manager, new_tensors_device, cache_indirection_buffer = self.sampler.get_cache_indirection() # Run model forward on the execution stream for proper synchronization - # with KVCacheTransferManager's onboard/offload operations. - self.execution_stream.wait_stream(torch.cuda.current_stream()) + # with KVCacheTransferManager's onboard/offload operations. With the + # loop on the execution stream both waits are program order. + if not self._loop_on_execution_stream: + self.execution_stream.wait_stream(torch.cuda.current_stream()) self._attach_encoder_output_to_execution_stream(scheduled_requests) with torch.cuda.stream(self.execution_stream): outputs = forward(scheduled_requests, self.resource_manager, @@ -8417,7 +8460,8 @@ def forward(scheduled_requests, resource_manager, new_tensors_device, # Ensure the default stream waits for execution_stream to complete # before downstream operations use the outputs. - torch.cuda.current_stream().wait_stream(self.execution_stream) + if not self._loop_on_execution_stream: + torch.cuda.current_stream().wait_stream(self.execution_stream) self._kv_connector_wait_for_save() if route_capture is not None: @@ -8546,9 +8590,11 @@ def _sample_async(self, scheduled_batch, beam_width = self.sampler.beam_width( scheduled_batch.all_requests()) + # A forward that sampled greedy tokens in place of the logits + # returns no logits; no request of such a batch asks for them. HandleLogits()(scheduled_batch.context_requests, scheduled_batch.generation_requests, - batch_outputs["logits"], beam_width, + batch_outputs.get("logits"), beam_width, num_context_logits_prefix_sum, self.sampler.is_generation_model()) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 348328d75d96..84ef85352cf6 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -971,6 +971,7 @@ def allocation_scope(current_stage: ExecutorMemoryType): ): # GuidedDecoder is applicable to non-speculative decoding and two-model speculative decoding. guided_decoder = GuidedDecoder(**kwargs) + model_engine.drop_greedy_only_outputs() elif spec_config.spec_dec_mode.support_capturable_guided_decoder( ): # CapturableGuidedDecoder is applicable to one-model speculative decoding. diff --git a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py new file mode 100644 index 000000000000..8fa36c484922 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Step-loop behaviour a model may opt into, read once by the engine. + +A model declares its hooks as ``runtime_hooks`` (an instance attribute or a class +attribute). Every field defaults to off, and a model that declares nothing gets +``RuntimeHooks()``: the step loop, the samplers and the attention metadata then +run exactly as they do without this module. + +The hooks change *where* and *how* the step loop moves data, never what a +forward computes. Each one is armed only by a model that has measured it on its +own deployment; none is a general default. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional, Protocol + +if TYPE_CHECKING: + import torch + + +class SpecStepFold(Protocol): + """Speculative step-input staging folded into the previous step's forward. + + A one-model speculative step ends by storing its accepted tokens, draft + tokens and lengths into the sampler's per-slot store, and the next step + begins by gathering those values back into the engine's flat input buffers. + A model that implements this protocol does both inside its last draft pass: + it writes the sampler store and stages the next step's inputs directly, and + the engine and the sampler skip their copies on the steps where it did. + + The object is owned by the model; the engine and the sampler only call it. + Every method is host-side bookkeeping or a single kernel launch on the + current stream. + """ + + 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: + """Receive the engine's persistent speculative input buffers, once, at construction.""" + + def bind_sampler_store(self, store: object) -> None: + """Receive the speculative sampler's persistent per-slot store, once, at construction.""" + + def begin_step(self, seq_slots: Optional[list[int]], num_rows: int) -> bool: + """Start a step of ``num_rows`` generation rows. + + ``seq_slots`` are those rows' sequence slots on a decode-only step whose + rows all have one, else None: the forward then stores nothing for them. + Returns whether the previous step's forward staged its inputs for + ``seq_slots``; the engine additionally checks that this step adds no + inputs of its own before it skips its gathers. + """ + + 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: + """Do the engine's previous-batch gathers in one launch. + + Writes the engine buffers exactly as the engine's own gathers would, + including zeroing both offset buffers outside the gathered range. + Returns False, having written nothing, when it cannot take the step. + """ + + def sampler_rows_stored(self, outputs: dict, num_skip: int, seq_slots: list[int]) -> bool: + """Whether the forward already stored this step's outputs for ``seq_slots`` into the sampler store.""" + + 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 the step's outputs into the sampler store in one launch. + + Returns False, having written nothing, when it cannot take the step. + """ + + +class MsaKvUpdate(Protocol): + """Advance the MSA decode kernels' per-request state after the KV lengths moved. + + Called by ``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` with the + metadata's own buffers; recomputes the attended lengths, the new tokens' + cache slots and their valid-block counts in one launch, exactly as the + metadata's own update does. Returns False, having written nothing, when it + cannot take the step. + """ + + 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"], + ) -> bool: ... + + +@dataclass(frozen=True) +class MsaRuntimePolicy: + """MiniMax-M3 MSA metadata preparation a model may opt into. + + Every option changes how a step's selection and cache-write tables reach + the device, not which blocks are attended. + """ + + # A step whose every query token has at most top-k valid blocks attends all + # of them, so its selection is the identity table: prepare stages that table + # instead of a scoring plan, and decode CUDA graphs are keyed by short and + # long contexts. + short_sequence_identity: bool = False + # The model's one-launch replacement for the update after the KV lengths move. + kv_update: Optional[MsaKvUpdate] = None + # Skip a block-table copy whose inputs equal the previous step's. + 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 + + +@dataclass(frozen=True) +class RuntimeHooks: + """The step-loop hooks a model opts into. All off by default.""" + + # 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 + # Skip host-to-device metadata copies whose values the destination already holds. + skip_unchanged_h2d: bool = False + # 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 + # 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 + # Speculative step-input staging folded into the forward (see SpecStepFold). + spec_step_fold: Optional[SpecStepFold] = None + # MiniMax-M3 MSA metadata preparation. + msa: MsaRuntimePolicy = MsaRuntimePolicy() + + @property + def skips_unchanged_metadata_uploads(self) -> bool: + """Whether a hook skips uploads into graph metadata buffers that already hold the values. + + Such a skip is sound only while no other engine writes those buffers, so the engine + then gives its CUDA graph metadata a buffer pool of its own. + """ + return self.skip_unchanged_h2d or self.msa.skip_unchanged_block_offsets + + +DEFAULT_RUNTIME_HOOKS = RuntimeHooks() + + +def runtime_hooks_of(model: object) -> RuntimeHooks: + """The hooks ``model`` declares, or the all-off default when it declares none. + + Declaring hooks is optional for every model class, so absence is the common + case and means the default. + """ + hooks = getattr(model, "runtime_hooks", None) + if hooks is None: + return DEFAULT_RUNTIME_HOOKS + if not isinstance(hooks, RuntimeHooks): + raise TypeError( + f"{type(model).__name__}.runtime_hooks must be a RuntimeHooks, got {type(hooks).__name__}" + ) + return hooks diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 2464827a1a9a..2e2517c259a2 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -145,6 +145,34 @@ class SampleState(Generic[GenericSampleStateTensorsHost, GenericSampleStateTenso # Generic bounds not supported, https://github.com/python/typing/issues/548 GenericSampleState = TypeVar("GenericSampleState", bound=SampleState) # type: ignore +# Model output keys of a forward that samples greedily in place of its LM head: per-row greedy token ids, and the +# head's input rows so the logits can still be computed for a batch that needs them. Such a forward returns no +# logits. +GREEDY_TOKENS_KEY = "greedy_tokens" +GREEDY_HIDDEN_STATES_KEY = "greedy_hidden_states" + + +def greedy_tokens_usable(request: LlmRequest) -> bool: + """Whether a generation request is served by the forward's greedy token alone, without logits. + + The request conditions of TorchSampler's stable greedy path, plus every other reader of the logits. A dummy + request is not usable: the stable greedy path takes no batch that carries one. + """ + return not request.is_dummy and ( + get_draft_token_length(request) == 0 + and request.py_embedding_bias is None + and not getattr(request, "py_bad_words", None) + and not getattr(request, "py_no_repeat_ngram_size", None) + and not has_occurrence_penalty(request) + and not request.py_min_length + and not request.py_return_log_probs + and not request.py_return_generation_logits + and not request.py_logits_post_processors + and request.py_guided_decoding_params is None + and not check_stop_words_length(request) + and _request_strategy(request, vocab_size=2**31) == GREEDY + ) + def _fast_strategy_params(strategy: Strategy) -> tuple[float, int, float, float]: """Flatten a fast-tier strategy into per-row ``(temperature, top_k, top_p, min_p)``.""" @@ -2513,7 +2541,10 @@ def _process_requests( torch.Tensor, bool, ]: - raw_logits_cuda = model_outputs["logits"] + # A forward that sampled greedy tokens in place of the logits (GREEDY_TOKENS_KEY) may return no logits. + raw_logits_cuda = model_outputs.get("logits") + greedy_tokens = model_outputs.get(GREEDY_TOKENS_KEY) + rows_cuda = raw_logits_cuda if greedy_tokens is None else greedy_tokens generation_requests = scheduled_requests.generation_requests request_ids = [request.py_request_id for request in generation_requests] @@ -2526,7 +2557,7 @@ def _process_requests( bool(generation_requests) and self.max_beam_width == 1 and scheduled_requests.num_context_requests == 0 - and len(generation_requests) <= raw_logits_cuda.shape[0] + and len(generation_requests) <= rows_cuda.shape[0] and model_outputs.get("d2t") is None and all( not request.is_dummy and get_draft_token_length(request) == 0 @@ -2563,6 +2594,9 @@ def _process_requests( live_requests = [request for request in generation_requests if not request.is_dummy] staged_matches_batch = ( self._in_graph_staged + # A forward that sampled greedily in place of its LM head skipped the + # in-graph sampling, so the staged buffers were never written. + and greedy_tokens is None and self._in_graph_dest_indices is not None # A draft request reuses its target's py_request_id, so the id # list alone cannot tell a drafter batch from the target batch @@ -2660,13 +2694,19 @@ def _process_requests( self._stable_greedy_seq_slots_host = seq_slots_host self._stable_greedy_seq_slots_cuda = seq_slots_cuda - next_tokens = fast_greedy_sample_kernel( - raw_logits_cuda[: len(generation_requests)], - new_tokens_cuda, - seq_slots_cuda, - self.max_beam_width, - None, - ) + if greedy_tokens is not None: + next_tokens = greedy_tokens[: len(generation_requests)].to(new_tokens_cuda.dtype) + scatter_new_tokens( + next_tokens, new_tokens_cuda, seq_slots_cuda, self.max_beam_width + ) + else: + next_tokens = fast_greedy_sample_kernel( + raw_logits_cuda[: len(generation_requests)], + new_tokens_cuda, + seq_slots_cuda, + self.max_beam_width, + None, + ) new_tokens_host = self._copy_to_host(next_tokens) return ( generation_requests, @@ -2681,6 +2721,10 @@ def _process_requests( self._stable_greedy_request_ids = [] self._stable_greedy_seq_slots = [] + if raw_logits_cuda is None: + raise RuntimeError( + "the forward returned greedy tokens without logits for a batch that needs logits" + ) sampling_requests, sampling_requests_metadata, logits_cuda = self._select_generated_logits( scheduled_requests, raw_logits_cuda, diff --git a/tensorrt_llm/_torch/pyexecutor/step_inputs.py b/tensorrt_llm/_torch/pyexecutor/step_inputs.py new file mode 100644 index 000000000000..2f47dadaba8d --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/step_inputs.py @@ -0,0 +1,194 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pack a decode step's small per-step input copies into one upload and one scatter launch. + +Between two decode-graph replays, the next step's metadata reaches the device as a chain of tiny copies (positions, +KV lengths, cache slots, and the graph runner's static-input copies). Each is a separate operation on the stream the +next graph waits for, so the GPU idles for the whole chain. + +While a pack is open (``begin`` ... ``flush``), ``copy_(dst, src)`` records a host-to-device copy by snapshotting the +host bytes into a pinned staging buffer, and a device-to-device copy as a (dst, src) descriptor. ``flush`` issues one +host-to-device copy of [descriptors | data] into a device staging buffer and one ``trtllm::step_inputs_scatter`` +launch, which writes every staged host copy and then every device copy (in that order, so a device copy may read a +destination a host copy just filled). With no pack open, ``copy_`` is ``dst.copy_(src, non_blocking=True)``. + +Each engine that armed ``RuntimeHooks.pack_step_inputs`` owns one ``StepInputPack`` (on its device) and makes it the +active pack for its forward (``active``); the module functions act on the active pack of the calling context, so +executors on other threads or devices never share one. The engine opens it only when the step will reach the flush (a +replay of an already captured graph). Route a copy through ``copy_`` only when no device work reads or writes its +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. +""" + +from __future__ import annotations + +import contextlib +from contextvars import ContextVar +from typing import Iterator, Optional + +import torch + +from tensorrt_llm._utils import prefer_pinned + + +class StepInputPack: + """One engine's pack of deferred step-input copies (see the module docstring).""" + + CAPACITY = 1 << 16 + MAX_DESC = 256 + # Pinned staging buffers in rotation: a forced flush reopens the pack without waiting on the upload just issued. + RING = 4 + HEADER_BYTES = 16 + DESC_BYTES = 24 + ALIGN = 16 + + def __init__(self, device: torch.device) -> None: + self.host = [ + torch.empty(self.CAPACITY, dtype=torch.uint8, pin_memory=prefer_pinned()) + for _ in range(self.RING) + ] + self.events: list[Optional[torch.cuda.Event]] = [None] * self.RING + self.dev = torch.empty(self.CAPACITY, dtype=torch.uint8, device=device) + self.data_start = self.HEADER_BYTES + self.DESC_BYTES * self.MAX_DESC + self.parity = 0 + self.open = False + self.h2d: list[tuple[int, int, int]] = [] + 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.used = self.data_start + + def begin(self) -> None: + if self.open: + self.flush() + event = self.events[self.parity] + if event is not None: + # The upload that last read this pinned buffer has run. + event.synchronize() + self._clear() + self.open = True + + def _clear(self) -> None: + self.h2d.clear() + self.d2d.clear() + self.staged.clear() + self.used = self.data_start + + def stage(self, dst: torch.Tensor, src: torch.Tensor) -> bool: + if not dst.is_cuda or dst.numel() != src.numel(): + return False + nbytes = dst.numel() * dst.element_size() + if src.device.type == "cpu": + # A contiguous destination is one descriptor; a 2-D slice with contiguous rows (a page table's leading + # rows and columns) is one per row. + if dst.is_contiguous(): + rows, row_bytes, row_stride = 1, nbytes, 0 + elif dst.dim() == 2 and dst.stride(1) == 1: + rows = dst.shape[0] + row_bytes = dst.shape[1] * dst.element_size() + row_stride = dst.stride(0) * dst.element_size() + else: + return False + offset = (self.used + self.ALIGN - 1) // self.ALIGN * self.ALIGN + if ( + offset + nbytes > self.CAPACITY + or len(self.h2d) + len(self.d2d) + rows > self.MAX_DESC + ): + return False + # A dtype change converts on the host, as copy_ would on the device. + src = src.to(dst.dtype).contiguous() + self.host[self.parity][offset : offset + nbytes].copy_(src.view(-1).view(torch.uint8)) + base = dst.data_ptr() + for row in range(rows): + self.h2d.append((base + row * row_stride, offset + row * row_bytes, row_bytes)) + self.staged.append(dst) + self.used = offset + nbytes + return True + if len(self.h2d) + len(self.d2d) >= self.MAX_DESC: + return False + if src.is_cuda and src.dtype == dst.dtype and src.is_contiguous() and dst.is_contiguous(): + self.d2d.append((dst.data_ptr(), src.data_ptr(), nbytes)) + self.staged += (dst, src) + return True + return False + + def flush(self) -> None: + if not self.open: + return + self.open = False + if not self.h2d and not self.d2d: + return + host = self.host[self.parity] + desc = [len(self.h2d), len(self.d2d)] + [ + value for copy in self.h2d + self.d2d for value in copy + ] + 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) + event = self.events[self.parity] + if event is 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._clear() + + def discard(self) -> None: + """Close the pack without issuing it (its step failed).""" + self.open = False + self._clear() + + +_active: ContextVar[Optional[StepInputPack]] = ContextVar("trtllm_step_input_pack", default=None) + + +@contextlib.contextmanager +def active(pack: Optional[StepInputPack]) -> Iterator[None]: + """Make ``pack`` the calling context's pack (None: no packing) for the duration.""" + token = _active.set(pack) + try: + yield + finally: + _active.reset(token) + + +def begin() -> None: + """Open the active pack for the coming copies. Call only when the step will reach ``flush``.""" + pack = _active.get() + if pack is not None: + pack.begin() + + +def is_open() -> bool: + pack = _active.get() + return pack is not None and pack.open + + +def copy_(dst: torch.Tensor, src: torch.Tensor) -> torch.Tensor: + """``dst.copy_(src, non_blocking=True)``, deferred into the active open pack when there is one.""" + pack = _active.get() + if pack is not None and pack.open and not torch.cuda.is_current_stream_capturing(): + if pack.stage(dst, src): + return dst + # Keep the order: everything staged so far lands before this copy, and the copies after it pack again. + pack.flush() + out = dst.copy_(src, non_blocking=True) + pack.begin() + return out + return dst.copy_(src, non_blocking=True) + + +def discard() -> None: + """Close the active pack without issuing it (its step failed).""" + pack = _active.get() + if pack is not None: + pack.discard() + + +def flush() -> None: + """Issue the active open pack (one upload and one scatter launch on the current stream) and close it.""" + pack = _active.get() + if pack is not None: + pack.flush() diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index fc29a0633f41..c766c5027624 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -17,6 +17,7 @@ from ..attention.backends.flashinfer import FlashInferAttentionMetadata from ..attention.backends.sparse.params import MTPIndexShareMetadata from ..model_config import ModelConfig +from ..pyexecutor.h2d_cache import copy_h2d_if_changed from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.llm_request import LlmRequest from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager @@ -352,8 +353,7 @@ def prepare(self): dtype=torch.int, device='cpu', pin_memory=prefer_pinned()) - self.batch_indices_cuda[:num_seqs].copy_(batch_indices, - non_blocking=True) + copy_h2d_if_changed(self.batch_indices_cuda[:num_seqs], batch_indices) # Snapshot the real (unpadded) scheduled token count before the decrement # below. maybe_capture_hidden_states runs during the target forward and # must bound the capture by this count; self.num_tokens is about to be @@ -374,8 +374,7 @@ def prepare(self): eagle_slot_ids_tensor = torch.tensor(eagle_slot_ids, dtype=torch.int, pin_memory=prefer_pinned()) - self.slot_ids[:num_seqs].copy_(eagle_slot_ids_tensor, - non_blocking=True) + copy_h2d_if_changed(self.slot_ids[:num_seqs], eagle_slot_ids_tensor) sa_manager = getattr(self.spec_resource_manager, 'sa_manager', None) if sa_manager is not None: diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 724b2aa15e36..596d644315d9 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -584,6 +584,10 @@ def is_ragged_verify(self) -> bool: group_all_greedy_sample: Optional[bool] = None # Whether to use rejection sampling for one-model speculative decoding. use_rejection_sampling: bool = False + # Skip the Philox seed and offset uploads on all-greedy batches (deploy-time; + # 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 # 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 @@ -899,6 +903,9 @@ def _upload(dst: torch.Tensor, values: list[int]) -> None: pin_memory=prefer_pinned()), non_blocking=True) + if self.skip_greedy_rng_upload and self.is_all_greedy_sample: + # The offset window counters above still advanced. + return _upload(self.request_seeds, request_seeds) _upload(self.request_offsets, request_offsets) _upload(self.seeds, flat_seeds) diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 5fb5b185dc0b..8a51b3e69a35 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -30,6 +30,7 @@ from ...sampling_params import SamplingParams from ..pyexecutor.llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..pyexecutor.resource_manager import BaseResourceManager +from ..pyexecutor.runtime_hooks import SpecStepFold from ..pyexecutor.sampler import ( DEFAULT_BEAM_IDX, AsyncWorkerMixin, @@ -165,6 +166,9 @@ class SpecSampler(Sampler[SampleStateSpec], AsyncWorkerMixin): draft length every mode used to compute for itself. """ + # The model's speculative step-input staging, when it opted into one (bind_spec_step_fold). + spec_step_fold: Optional[SpecStepFold] = None + SampleState = SampleStateSpec def is_generation_model(self) -> bool: @@ -313,6 +317,11 @@ def __init__( new_tokens_lens=int_tensor((seq_slots,)), ) + 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 + spec_step_fold.bind_sampler_store(self.store) + def _request_common_handling( self, request: LlmRequest, @@ -396,6 +405,54 @@ def update_requests( req.py_rewind_len = verified_len - req.py_num_accepted_draft_tokens self._request_common_handling(req, next_draft_tokens_list, runtime_draft_len) + def _store_outputs( + self, + seq_slots: list[int], + o_new_tokens: torch.Tensor, + o_next_new_tokens: torch.Tensor, + o_new_tokens_lens: torch.Tensor, + o_next_draft_tokens: torch.Tensor, + ) -> None: + """Store one step's worker outputs into the per-slot store.""" + slots = torch.as_tensor(seq_slots, dtype=torch.long) + slots = slots.to(device="cuda", non_blocking=True) + step_fold = self.spec_step_fold + if step_fold is not None and step_fold.scatter( + slots, o_new_tokens, o_next_new_tokens, o_new_tokens_lens, o_next_draft_tokens + ): + return + + # Pad or truncate to match fixed-size store buffers for index_copy_. + # The worker output width tracks runtime_draft_len, which dynamic draft + # length shrinks below the statically allocated store width. + new_tokens_width = self.store.new_tokens.shape[0] + next_new_tokens_width = self.store.next_new_tokens.shape[0] + draft_tokens_width = self.store.next_draft_tokens.shape[1] + if o_new_tokens.shape[1] < new_tokens_width: + o_new_tokens = torch.nn.functional.pad( + o_new_tokens, (0, new_tokens_width - o_new_tokens.shape[1]) + ) + elif o_new_tokens.shape[1] > new_tokens_width: + o_new_tokens = o_new_tokens[:, :new_tokens_width] + if o_next_draft_tokens.shape[1] < draft_tokens_width: + o_next_draft_tokens = torch.nn.functional.pad( + o_next_draft_tokens, (0, draft_tokens_width - o_next_draft_tokens.shape[1]) + ) + elif o_next_draft_tokens.shape[1] > draft_tokens_width: + o_next_draft_tokens = o_next_draft_tokens[:, :draft_tokens_width] + if o_next_new_tokens.shape[1] < next_new_tokens_width: + o_next_new_tokens = torch.nn.functional.pad( + o_next_new_tokens, (0, next_new_tokens_width - o_next_new_tokens.shape[1]) + ) + elif o_next_new_tokens.shape[1] > next_new_tokens_width: + o_next_new_tokens = o_next_new_tokens[:, :next_new_tokens_width] + + # Use index_copy_ for efficient copying (slots are unique) + self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens) + self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens) + self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens) + self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens) + def sample_async( self, scheduled_requests: ScheduledRequests, @@ -442,8 +499,7 @@ def sample_async( for r in sampling_requests ] - slots = torch.as_tensor([r.py_seq_slot for r in sampling_requests], dtype=torch.long) - slots = slots.to(device="cuda", non_blocking=True) + seq_slots = [r.py_seq_slot for r in sampling_requests] 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] @@ -489,36 +545,11 @@ def sample_async( raise ValueError("verify_lens must contain one int32 token window per output row") o_verify_lens = o_verify_lens[num_skip : num_skip + num_sampling_requests] - # Pad or truncate to match fixed-size store buffers for index_copy_. - # The worker output width tracks runtime_draft_len, which dynamic draft - # length shrinks below the statically allocated store width. - new_tokens_width = self.store.new_tokens.shape[0] - next_new_tokens_width = self.store.next_new_tokens.shape[0] - draft_tokens_width = self.store.next_draft_tokens.shape[1] - if o_new_tokens.shape[1] < new_tokens_width: - o_new_tokens = torch.nn.functional.pad( - o_new_tokens, (0, new_tokens_width - o_new_tokens.shape[1]) - ) - elif o_new_tokens.shape[1] > new_tokens_width: - o_new_tokens = o_new_tokens[:, :new_tokens_width] - if o_next_draft_tokens.shape[1] < draft_tokens_width: - o_next_draft_tokens = torch.nn.functional.pad( - o_next_draft_tokens, (0, draft_tokens_width - o_next_draft_tokens.shape[1]) - ) - elif o_next_draft_tokens.shape[1] > draft_tokens_width: - o_next_draft_tokens = o_next_draft_tokens[:, :draft_tokens_width] - if o_next_new_tokens.shape[1] < next_new_tokens_width: - o_next_new_tokens = torch.nn.functional.pad( - o_next_new_tokens, (0, next_new_tokens_width - o_next_new_tokens.shape[1]) + step_fold = self.spec_step_fold + if step_fold is None or not step_fold.sampler_rows_stored(outputs, num_skip, seq_slots): + self._store_outputs( + seq_slots, o_new_tokens, o_next_new_tokens, o_new_tokens_lens, o_next_draft_tokens ) - elif o_next_new_tokens.shape[1] > next_new_tokens_width: - o_next_new_tokens = o_next_new_tokens[:, :next_new_tokens_width] - - # Use index_copy_ for efficient copying (slots are unique) - self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens) - self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens) - self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens) - self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens) # Create sample state with async D2H copy device_tensors = SampleStateTensorsSpec( diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py index c9945d39d48c..61026b6a5f65 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py @@ -175,11 +175,13 @@ def test_index_decode_score_matches_reference(dtype, num_heads, decode_query_len @skip_not_sm100 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) -def test_index_decode_score_multi_round_split_k(dtype): +@pytest.mark.parametrize("split_k", [None, 256, 512], ids=["default", "split256", "split512"]) +def test_index_decode_score_multi_round_split_k(dtype, split_k): """A request longer than split_k pages makes the per-CTA block loop iterate. - split_k is 256, so 33000 tokens is 258 blocks and CTAs 0 and 1 each handle - two blocks while the rest handle one. + The default split_k is 256, so 33000 tokens is 258 blocks and CTAs 0 and 1 + each handle two blocks while the rest handle one. With 512 every block has + its own CTA and the trailing CTAs have none. """ seq_lens = [33000, 32768 + 1] idx_q, k_cache, block_table, seq_lens_dev, _, score, expected = _make_inputs( @@ -187,12 +189,39 @@ def test_index_decode_score_multi_round_split_k(dtype): ) assert (max(seq_lens) + PAGE_SIZE - 1) // PAGE_SIZE > 256 + split_k_args = () if split_k is None else (split_k,) torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( - idx_q, k_cache, block_table, seq_lens_dev, score, 1 + idx_q, k_cache, block_table, seq_lens_dev, score, 1, *split_k_args ) _assert_scores_close(score, expected) +@skip_not_sm100 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize( + ("num_heads", "decode_query_len", "seq_lens"), + [(4, 1, [300, 1500, 2049, 33000]), (4, 4, [4097, 70000])], + ids=["decode", "spec_verify"], +) +def test_index_decode_score_does_not_depend_on_split_k( + dtype, num_heads, decode_query_len, seq_lens +): + """Each block is scored by exactly one CTA, so the scores are identical across split_k.""" + idx_q, k_cache, block_table, seq_lens_dev, _, score, expected = _make_inputs( + seq_lens, dtype=dtype, num_heads=num_heads, decode_query_len=decode_query_len, seed=5 + ) + scores = {} + for split_k in (256, 512): + score.fill_(-float("inf")) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, k_cache, block_table, seq_lens_dev, score, decode_query_len, split_k + ) + scores[split_k] = score.clone() + + assert torch.equal(scores[256], scores[512]) + _assert_scores_close(scores[512], expected) + + @skip_not_sm100 def test_index_decode_score_transposed_view_matches_contiguous(): """The transposed selector view must produce the same values as a direct write. 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 8c83072141c8..05e053be775d 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -42,6 +42,7 @@ ) from tensorrt_llm._torch.attention.backends.sparse.registry import _resolve_minimax_m3_backend_cls from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy from tensorrt_llm.bindings import DataType from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig @@ -803,6 +804,7 @@ class FakeMetadata: # indexer takes the proxy path and needs no decode span. msa_decode_span = None msa_prefill_proxy_plan = (False, 0, 2, {}, None) + _msa_prefill_identity = None msa_prefill_n_valid_blocks = torch.ones(2, dtype=torch.int32, device="cuda") msa_kv_indices = torch.arange(2, dtype=torch.int32, device="cuda") msa_qo_lens_cpu = torch.ones(2, dtype=torch.int32) @@ -920,6 +922,8 @@ def select_blocks(self, *args: object, **kwargs: object) -> torch.Tensor: return torch.zeros(num_tokens, 1, 16, dtype=torch.int32) class FakeMetadata: + _msa_prefill_identity = None + msa_prefill_n_valid_blocks = torch.ones(num_tokens, dtype=torch.int32) msa_n_valid_blocks = torch.ones(num_tokens, dtype=torch.int32) msa_worst_case_max_k_tiles = 8 @@ -1949,6 +1953,122 @@ def test_on_update_kv_lens_is_a_noop_without_speculative_decoding(monkeypatch): assert metadata.msa_n_valid_blocks.tolist() == [0] * 5 +class _RecordingKvUpdate: + """A model's one-launch KV update (MsaKvUpdate) that takes or declines every step.""" + + def __init__(self, taken: bool) -> None: + self._taken = taken + self.calls = [] + + def __call__( + self, + 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, + ) -> bool: + self.calls.append( + dict( + kv_lens=kv_lens, + kv_lens_staged=kv_lens_staged, + seq_lens=seq_lens, + q_batch_row=q_batch_row, + q_intra=q_intra, + block_table=block_table, + tokens_per_block=tokens_per_block, + batch=batch, + total_q=total_q, + msa_seq_lens=msa_seq_lens, + cu_kv_lens=cu_kv_lens, + out_cache_loc=out_cache_loc, + n_valid_blocks=n_valid_blocks, + ) + ) + return self._taken + + +def test_default_msa_policy_has_no_kv_update(monkeypatch): + metadata = _kv_lens_update_metadata( + monkeypatch, qo_lens=(1,), kv_staged=(5,), kv_corrected=(5,), page_size=4 + ) + + assert metadata.msa_policy == MsaRuntimePolicy() + assert metadata.msa_policy.kv_update is None + + +@pytest.mark.parametrize("prefill_counts", [False, True]) +def test_a_taken_kv_update_replaces_the_rederivation(monkeypatch, prefill_counts): + metadata = _kv_lens_update_metadata( + monkeypatch, + qo_lens=(2, 3), + kv_staged=(9, 12), + kv_corrected=(8, 12), + page_size=4, + prefill_counts=prefill_counts, + ) + kv_update = _RecordingKvUpdate(taken=True) + metadata.msa_policy = MsaRuntimePolicy(kv_update=kv_update) + + metadata.on_update_kv_lens() + + # The hook owns the update: the metadata's own re-derivation does not run. + assert metadata.msa_seq_lens_cuda.tolist() == [9, 12] + assert metadata.msa_out_cache_loc.tolist() == [-1] * 8 + assert metadata.msa_n_valid_blocks.tolist() == [0] * 8 + assert len(kv_update.calls) == 1 + call = kv_update.calls[0] + assert call["kv_lens"] is metadata.kv_lens_cuda + assert call["kv_lens_staged"] is metadata.msa_kv_lens_staged + assert call["q_batch_row"] is metadata.msa_q_batch_row + assert call["q_intra"] is metadata.msa_q_intra + assert call["block_table"] is metadata.msa_block_table + assert call["msa_seq_lens"] is metadata.msa_seq_lens_cuda + assert call["out_cache_loc"] is metadata.msa_out_cache_loc + assert call["seq_lens"].tolist() == [2, 3] + assert (call["tokens_per_block"], call["batch"], call["total_q"]) == (4, 2, 5) + assert call["cu_kv_lens"] is None + counts = metadata._msa_prefill_n_valid_blocks if prefill_counts else metadata.msa_n_valid_blocks + assert call["n_valid_blocks"].data_ptr() == counts.data_ptr() + assert call["n_valid_blocks"].numel() == 5 + + +def test_a_declined_kv_update_falls_back_to_the_rederivation(monkeypatch): + metadata = _kv_lens_update_metadata( + monkeypatch, qo_lens=(2, 3), kv_staged=(9, 12), kv_corrected=(8, 12), page_size=4 + ) + kv_update = _RecordingKvUpdate(taken=False) + metadata.msa_policy = MsaRuntimePolicy(kv_update=kv_update) + + metadata.on_update_kv_lens() + + assert len(kv_update.calls) == 1 + assert metadata.msa_seq_lens_cuda.tolist() == [8, 12] + assert metadata.msa_out_cache_loc[:5].tolist() == [6, 7, 109, 110, 111] + assert metadata.msa_n_valid_blocks[:5].tolist() == [2, 2, 3, 3, 3] + + +def test_kv_update_is_not_called_without_dynamic_kv_lens(monkeypatch): + metadata = _kv_lens_update_metadata( + monkeypatch, qo_lens=(1, 1), kv_staged=(4, 6), kv_corrected=(3, 5), page_size=4 + ) + metadata._msa_kv_lens_dynamic = False + kv_update = _RecordingKvUpdate(taken=True) + metadata.msa_policy = MsaRuntimePolicy(kv_update=kv_update) + + metadata.on_update_kv_lens() + + assert kv_update.calls == [] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize( ("src_dtype", "cache_dtype"), 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 new file mode 100644 index 000000000000..01c22219d453 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_runtime_policy.py @@ -0,0 +1,362 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 MSA decode-step preparation a model opts into (RuntimeHooks.msa). + +Covers the short-sequence identity table the metadata stages, and the KV cache +manager's block-table and pool-view shortcuts. With the default policy every +path runs as before; the existing test_msa_backend.py tests cover that. +""" + +from unittest.mock import create_autospec + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3KVCacheManagerV2 +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.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, + _fill_unused_slots, +) +from tensorrt_llm._torch.pyexecutor import h2d_cache +from tensorrt_llm._torch.pyexecutor.h2d_cache import H2DSkipCache +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy +from tensorrt_llm.bindings import DataType +from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig +from tensorrt_llm.mapping import Mapping + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + +NUM_KV_HEADS = 2 +PARAMS = MiniMaxM3SparseMetadataParams( + global_num_q_heads=8, global_num_kv_heads=NUM_KV_HEADS, num_index_heads=4, topk=16 +) + + +@pytest.fixture +def manager(monkeypatch) -> MiniMaxM3KVCacheManagerV2: + """An M3 cache manager whose V2 base allocates nothing (as in test_msa_backend.py).""" + + def fake_base_init(self, *args, **kwargs) -> None: + del args, kwargs + self.is_disagg = False + self.dtype = DataType.BF16 + self.layer_offsets = {3: 0} + + def no_index_k_buffer(self, layer_idx, **kwargs): + del self, layer_idx, kwargs + return None + + 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) + return MiniMaxM3KVCacheManagerV2( + num_layers=4, sparse_attention_config=MiniMaxM3SparseAttentionConfig() + ) + + +# --------------------------------------------------------------------------- +# Short-sequence identity table +# --------------------------------------------------------------------------- + + +def _short_table_metadata( + policy: MsaRuntimePolicy, manager: MiniMaxM3KVCacheManagerV2, num_contexts: int = 0 +) -> MiniMaxM3MsaSparseAttentionMetadata: + metadata = object.__new__(MiniMaxM3MsaSparseAttentionMetadata) + metadata.msa_policy = policy + metadata._num_contexts = num_contexts + metadata.kv_cache_manager = manager + metadata.max_num_sequences = 4 + metadata.max_num_tokens = 16 + metadata.cuda_graph_buffers = None + metadata.is_cuda_graph = False + metadata.mapping = Mapping() + return metadata + + +def _identity(n_valid: list[int]) -> torch.Tensor: + ar = torch.arange(MSA_REQUIRED_TOPK, dtype=torch.int32) + valid = torch.tensor(n_valid, dtype=torch.int32) + table = torch.where(ar[None, :] < valid[:, None], ar[None, :], -1) + return table.unsqueeze(0).expand(NUM_KV_HEADS, len(n_valid), MSA_REQUIRED_TOPK) + + +def _staged(metadata: MiniMaxM3MsaSparseAttentionMetadata, total_q: int) -> torch.Tensor: + torch.cuda.synchronize() + size = NUM_KV_HEADS * total_q * MSA_REQUIRED_TOPK + return metadata.msa_short_table[:size].view(NUM_KV_HEADS, total_q, MSA_REQUIRED_TOPK).cpu() + + +@requires_cuda +@pytest.mark.parametrize("n_valid", [[1], [1, 16, 5], [16, 16, 16, 16]]) +def test_short_steps_stage_the_identity_table(manager, n_valid): + metadata = _short_table_metadata(MsaRuntimePolicy(short_sequence_identity=True), manager) + + metadata._stage_msa_short_table(torch.tensor(n_valid, dtype=torch.int32), PARAMS) + + assert metadata._msa_short_ok + assert metadata.msa_short_table.is_cuda + assert torch.equal(_staged(metadata, len(n_valid)), _identity(n_valid)) + + +@requires_cuda +def test_the_default_policy_stages_nothing(manager): + metadata = _short_table_metadata(MsaRuntimePolicy(), manager) + + metadata._stage_msa_short_table(torch.tensor([1, 2], dtype=torch.int32), PARAMS) + + assert not metadata._msa_short_ok + assert metadata.msa_short_table is None + + +@requires_cuda +@pytest.mark.parametrize( + ("n_valid", "num_contexts"), + [([1, MSA_REQUIRED_TOPK + 1], 0), ([], 0), ([1, 2], 1)], + ids=["more_than_topk_blocks", "empty_step", "context_step"], +) +def test_ineligible_steps_do_not_use_the_identity_table(manager, n_valid, num_contexts): + metadata = _short_table_metadata(MsaRuntimePolicy(short_sequence_identity=True), manager) + # A previous eligible step staged the table. + metadata._stage_msa_short_table(torch.tensor([3], dtype=torch.int32), PARAMS) + assert metadata._msa_short_ok + + metadata._num_contexts = num_contexts + metadata._stage_msa_short_table(torch.tensor(n_valid, dtype=torch.int32), PARAMS) + + assert not metadata._msa_short_ok + + +@requires_cuda +@pytest.mark.parametrize("skip_cache", [False, True], ids=["plain", "skip_cache"]) +def test_the_identity_table_follows_each_step(manager, skip_cache): + """Also with the engine's skip cache active, which skips the repeated [2, 4] upload.""" + metadata = _short_table_metadata(MsaRuntimePolicy(short_sequence_identity=True), manager) + + with h2d_cache.active(H2DSkipCache() if skip_cache else None): + for n_valid in ([2, 3], [2, 4], [2, 4], [16], [2, 3]): + metadata._stage_msa_short_table(torch.tensor(n_valid, dtype=torch.int32), PARAMS) + assert torch.equal(_staged(metadata, len(n_valid)), _identity(n_valid)) + + +@requires_cuda +def test_the_identity_table_grows_past_its_first_bound(manager): + metadata = _short_table_metadata(MsaRuntimePolicy(short_sequence_identity=True), manager) + metadata._stage_msa_short_table(torch.tensor([1], dtype=torch.int32), PARAMS) + first_numel = metadata.msa_short_table.numel() + n_valid = [i % MSA_REQUIRED_TOPK + 1 for i in range(first_numel // MSA_REQUIRED_TOPK)] + + metadata._stage_msa_short_table(torch.tensor(n_valid, dtype=torch.int32), PARAMS) + + assert metadata._msa_short_ok + assert metadata.msa_short_table.numel() >= NUM_KV_HEADS * len(n_valid) * MSA_REQUIRED_TOPK + assert torch.equal(_staged(metadata, len(n_valid)), _identity(n_valid)) + + +# --------------------------------------------------------------------------- +# Block-table copy shortcut +# --------------------------------------------------------------------------- + + +@pytest.fixture +def base_copy(monkeypatch): + """The V2 base copy, recorded instead of launched.""" + copy = create_autospec(KVCacheManagerV2.copy_batch_block_offsets, spec_set=True) + monkeypatch.setattr(KVCacheManagerV2, "copy_batch_block_offsets", copy) + return copy.mock + + +def _block_table_host(manager: MiniMaxM3KVCacheManagerV2) -> torch.Tensor: + manager.host_kv_cache_block_offsets = torch.arange(2 * 4 * 2 * 8, dtype=torch.int32).view( + 2, 4, 2, 8 + ) + return manager.host_kv_cache_block_offsets + + +def _decode_copy(manager, dst, request_ids=(0, 1), num_contexts=0, max_blocks=8): + manager.copy_batch_block_offsets( + dst, list(request_ids), 1, num_contexts, len(request_ids), max_blocks=max_blocks + ) + + +def test_block_offsets_shortcut_is_off_by_default(manager, base_copy): + _block_table_host(manager) + dst = torch.zeros(2, 4, 2, 8, dtype=torch.int32) + + _decode_copy(manager, dst) + _decode_copy(manager, dst) + + assert not manager.skip_unchanged_block_offsets + assert base_copy.call_count == 2 + + +def test_an_unchanged_decode_copy_is_skipped(manager, base_copy): + _block_table_host(manager) + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=True, reuse_pool_views=False) + dst = torch.zeros(2, 4, 2, 8, dtype=torch.int32) + + _decode_copy(manager, dst) + _decode_copy(manager, dst) + + assert base_copy.call_count == 1 + + +def test_a_changed_host_table_is_copied(manager, base_copy): + host = _block_table_host(manager) + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=True, reuse_pool_views=False) + dst = torch.zeros(2, 4, 2, 8, dtype=torch.int32) + + _decode_copy(manager, dst) + host[0, 1, 0, 3] += 100 # a request grew into a new page + _decode_copy(manager, dst) + _decode_copy(manager, dst) + + assert base_copy.call_count == 2 + + +@pytest.mark.parametrize( + "change", + [ + dict(request_ids=(1, 0)), + dict(request_ids=(0, 2)), + dict(max_blocks=4), + ], + ids=["batch_order", "request_set", "max_blocks"], +) +def test_a_changed_batch_is_copied(manager, base_copy, change): + _block_table_host(manager) + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=True, reuse_pool_views=False) + dst = torch.zeros(2, 4, 2, 8, dtype=torch.int32) + + _decode_copy(manager, dst) + _decode_copy(manager, dst, **change) + + assert base_copy.call_count == 2 + + +def test_each_destination_keeps_its_own_record(manager, base_copy): + _block_table_host(manager) + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=True, reuse_pool_views=False) + first = torch.zeros(2, 4, 2, 8, dtype=torch.int32) + second = torch.zeros(2, 4, 2, 8, dtype=torch.int32) + + _decode_copy(manager, first) + _decode_copy(manager, second) + _decode_copy(manager, first) + _decode_copy(manager, second) + + assert base_copy.call_count == 2 + + +def test_a_step_with_contexts_always_copies_and_resets_the_record(manager, base_copy): + _block_table_host(manager) + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=True, reuse_pool_views=False) + dst = torch.zeros(2, 4, 2, 8, dtype=torch.int32) + + _decode_copy(manager, dst) + _decode_copy(manager, dst, num_contexts=1) + _decode_copy(manager, dst, num_contexts=1) + # The context steps wrote dst through the base copy: the next decode copy must run. + _decode_copy(manager, dst) + + assert base_copy.call_count == 4 + + +# --------------------------------------------------------------------------- +# Pool-view shortcut +# --------------------------------------------------------------------------- + +NUM_SLOTS, SCALE, PAGE_SHAPE = 3, 4, [2, 8, 16] + + +@pytest.fixture +def pools(monkeypatch): + """Two pools the manager's slot geometry resolves to: one for K/V, one for block scales.""" + numel = NUM_SLOTS * SCALE * 2 * 8 * 16 + kv_pool = torch.zeros(numel, dtype=torch.bfloat16, device="cuda") + scale_pool = torch.zeros(numel, dtype=torch.uint8, device="cuda") + + def kv_slot_geometry(self, layer_idx, kv_layout): + del self, layer_idx, kv_layout + return kv_pool.data_ptr(), torch.bfloat16, NUM_SLOTS, SCALE, list(PAGE_SHAPE) + + def kv_scale_slot_geometry(self, layer_idx, kv_layout): + del self, layer_idx, kv_layout + return scale_pool.data_ptr(), torch.uint8, NUM_SLOTS, SCALE, list(PAGE_SHAPE) + + monkeypatch.setattr(MiniMaxM3KVCacheManagerV2, "_kv_slot_geometry", kv_slot_geometry) + monkeypatch.setattr( + MiniMaxM3KVCacheManagerV2, "_kv_scale_slot_geometry", kv_scale_slot_geometry + ) + return kv_pool, scale_pool + + +def _assert_kv_view(view: torch.Tensor, pool: torch.Tensor) -> None: + assert view.data_ptr() == pool.data_ptr() + assert list(view.shape) == [NUM_SLOTS, 2, *PAGE_SHAPE] + assert view.stride(0) == SCALE * 2 * 8 * 16 + + +@requires_cuda +def test_pool_views_are_rebuilt_on_every_call_by_default(manager, pools): + kv_pool, _ = pools + + first = manager.get_buffers(3) + second = manager.get_buffers(3) + + assert not manager.reuse_pool_views + assert first is not second + _assert_kv_view(first, kv_pool) + _assert_kv_view(second, kv_pool) + + +@requires_cuda +def test_reused_pool_views_are_built_once_per_geometry(manager, pools): + kv_pool, scale_pool = pools + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=False, reuse_pool_views=True) + + view = manager.get_buffers(3) + scale_view = manager.get_block_scale_buffers(3) + + assert manager.get_buffers(3) is view + assert manager.get_block_scale_buffers(3) is scale_view + assert scale_view is not view + _assert_kv_view(view, kv_pool) + assert scale_view.data_ptr() == scale_pool.data_ptr() + assert scale_view.dtype == torch.uint8 + + +@requires_cuda +def test_reused_pool_views_are_kept_per_inference_mode(manager, pools): + kv_pool, _ = pools + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=False, reuse_pool_views=True) + + view = manager.get_buffers(3) + with torch.inference_mode(): + inference_view = manager.get_buffers(3) + assert manager.get_buffers(3) is inference_view + + assert inference_view is not view + assert inference_view.is_inference() + assert not view.is_inference() + _assert_kv_view(inference_view, kv_pool) + + +@requires_cuda +def test_unused_slots_are_filled_after_a_full_buffer_step(): + buf = torch.full((8,), 7, dtype=torch.int32, device="cuda") + with h2d_cache.active(H2DSkipCache()): + _fill_unused_slots(buf, 3) + # A step whose live rows cover the whole buffer overwrites the filled tail. + buf.copy_(torch.arange(8, dtype=torch.int32, device="cuda")) + _fill_unused_slots(buf, buf.shape[0]) + buf[:3] = 0 + _fill_unused_slots(buf, 3) + assert buf.tolist() == [0, 0, 0, -1, -1, -1, -1, -1] diff --git a/tests/unittest/_torch/executor/engine/test_decoder_forward_scopes.py b/tests/unittest/_torch/executor/engine/test_decoder_forward_scopes.py new file mode 100644 index 000000000000..f244495bcf10 --- /dev/null +++ b/tests/unittest/_torch/executor/engine/test_decoder_forward_scopes.py @@ -0,0 +1,242 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DecoderRunner.forward scopes its own step-input pack and H2D skip cache to the step.""" + +from typing import Callable, Optional +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor import h2d_cache, step_inputs +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as decoder_runner_module +from tensorrt_llm._torch.pyexecutor.h2d_cache import H2DSkipCache, copy_h2d_if_changed +from tensorrt_llm._torch.pyexecutor.sampler.sampler import ( + GREEDY_HIDDEN_STATES_KEY, + GREEDY_TOKENS_KEY, +) +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +class _ModelWithoutWeights(torch.nn.Module): + """What DecoderRunner.forward reads from the model before it runs a step.""" + + extra_attrs: dict = {} + + +class _StepFailed(RuntimeError): + pass + + +def _runner( + step: Callable[..., object], + skip_cache: Optional[H2DSkipCache], + pack: Optional[StepInputPack], +) -> DecoderRunner: + # forward reads only the model's extra attributes, the two per-engine objects and the step. + runner = object.__new__(DecoderRunner) + runner.model = _ModelWithoutWeights() + runner._h2d_cache = skip_cache + runner._step_pack = pack + runner._forward_decoder = step + return runner + + +def _host(values): + return torch.tensor(values, dtype=torch.int32).pin_memory() + + +def test_forward_activates_the_runner_cache_and_pack_for_the_step_only(): + skip_cache = H2DSkipCache() + pack = StepInputPack(torch.device("cuda", torch.cuda.current_device())) + dst = torch.zeros(3, dtype=torch.int32, device="cuda") + values = _host([1, 2, 3]) + seen = {} + + def step(inputs, resource_manager, is_dummy=False): + seen["skip_enabled"] = h2d_cache.enabled() + step_inputs.begin() + seen["pack_open"] = pack.open + seen["copies"] = [copy_h2d_if_changed(dst, values) for _ in range(2)] + step_inputs.flush() + return "outputs" + + runner = _runner(step, skip_cache, pack) + + assert runner.forward(object(), resource_manager=None) == "outputs" + + assert seen == {"skip_enabled": True, "pack_open": True, "copies": [True, False]} + # Outside the step neither is active any more. + assert not h2d_cache.enabled() + assert not step_inputs.is_open() + # The shadow the step recorded belongs to the runner's own cache. + with h2d_cache.active(skip_cache): + assert copy_h2d_if_changed(dst, values) is False + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3] + + +def test_forward_without_armed_hooks_activates_nothing(): + seen = {} + + def step(inputs, resource_manager, is_dummy=False): + step_inputs.begin() + seen["skip_enabled"] = h2d_cache.enabled() + seen["pack_open"] = step_inputs.is_open() + + _runner(step, None, None).forward(object(), resource_manager=None) + + assert seen == {"skip_enabled": False, "pack_open": False} + + +def test_forward_hides_an_outer_context_cache_and_pack(): + """Another engine's objects, active around this call, are not used by this runner's step.""" + outer_cache = H2DSkipCache() + outer_pack = StepInputPack(torch.device("cuda", torch.cuda.current_device())) + seen = {} + + def step(inputs, resource_manager, is_dummy=False): + seen["skip_enabled"] = h2d_cache.enabled() + seen["pack_open"] = step_inputs.is_open() + + with h2d_cache.active(outer_cache), step_inputs.active(outer_pack): + outer_pack.begin() + _runner(step, None, None).forward(object(), resource_manager=None) + assert step_inputs.is_open() + outer_pack.discard() + + assert seen == {"skip_enabled": False, "pack_open": False} + + +def test_a_failed_step_discards_its_pack_and_clears_its_cache(): + skip_cache = H2DSkipCache() + pack = StepInputPack(torch.device("cuda", torch.cuda.current_device())) + dst = torch.zeros(3, dtype=torch.int32, device="cuda") + values = _host([7, 8, 9]) + + def failing_step(inputs, resource_manager, is_dummy=False): + step_inputs.begin() + assert copy_h2d_if_changed(dst, values, pack=True) is True + assert pack.staged + raise _StepFailed("forward failed after staging its inputs") + + runner = _runner(failing_step, skip_cache, pack) + + with pytest.raises(_StepFailed): + runner.forward(object(), resource_manager=None) + + assert not pack.open + assert pack.staged == [] and pack.h2d == [] and pack.d2d == [] + assert not h2d_cache.enabled() and not step_inputs.is_open() + torch.cuda.synchronize() + assert dst.tolist() == [0, 0, 0] + + # The next step re-sends the values the failed one never delivered. + def next_step(inputs, resource_manager, is_dummy=False): + step_inputs.begin() + sent = copy_h2d_if_changed(dst, values, pack=True) + step_inputs.flush() + return sent + + runner._forward_decoder = next_step + if hasattr(torch.ops.trtllm, "step_inputs_scatter"): + assert runner.forward(object(), resource_manager=None) is True + torch.cuda.synchronize() + assert dst.tolist() == [7, 8, 9] + else: + # Without the scatter op the pack cannot be issued; the cache decision is what matters. + with h2d_cache.active(skip_cache): + assert copy_h2d_if_changed(dst, values) is True + + +class _LogitsProcessor: + def forward(self, hidden_states, lm_head, attn_metadata, return_context_logits): + return lm_head(hidden_states) * 2 + + +class _GreedyFallbackModel(torch.nn.Module): + """The LM head and logits processor the greedy fallback runs eagerly.""" + + def __init__(self) -> None: + super().__init__() + self.lm_head = torch.nn.Identity() + self.logits_processor = _LogitsProcessor() + + +class _ScheduledWithContext: + """A batch with a context request: greedy tokens cannot serve it.""" + + num_context_requests = 1 + generation_requests = () + + +def test_greedy_fallback_leaves_the_replayed_graph_outputs_intact(): + runner = object.__new__(DecoderRunner) + runner.model = _GreedyFallbackModel() + # max_beam_width is read from the runner config; the fallback reads nothing else from it. + runner._config = Mock(spec_set=["max_beam_width"], max_beam_width=1) + runner.sample_in_graph_callable = None + hidden = torch.ones(2, 4, device="cuda") + tokens = torch.zeros(2, dtype=torch.int32, device="cuda") + # The dict a CUDA graph runner returns on every replay of its key. + graph_outputs = {GREEDY_TOKENS_KEY: tokens, GREEDY_HIDDEN_STATES_KEY: hidden} + + outputs = runner._logits_for_greedy_tokens(_ScheduledWithContext(), graph_outputs) + + assert GREEDY_TOKENS_KEY not in outputs and GREEDY_HIDDEN_STATES_KEY not in outputs + assert torch.equal(outputs["logits"], hidden * 2) + # The next replay of the same graph still sees its greedy outputs and takes the fallback again. + assert graph_outputs == {GREEDY_TOKENS_KEY: tokens, GREEDY_HIDDEN_STATES_KEY: hidden} + + +class _ScheduledGenerationOnly: + """A generation-only batch every request of which greedy tokens could serve (greedy_tokens_usable patched).""" + + num_context_requests = 0 + generation_requests = (object(),) + + +class _ScheduledIdle: + """An idle attention-DP rank's empty batch.""" + + num_context_requests = 0 + generation_requests = () + + +@pytest.mark.parametrize("keep", [True, False]) +def test_greedy_only_outputs_follow_the_guided_decoder_switch(keep, monkeypatch): + """With an executor-side guided decoder (keep_greedy_only_outputs False) the batch gets logits; otherwise the + forward's greedy-only outputs are returned as they are.""" + runner = object.__new__(DecoderRunner) + runner.model = _GreedyFallbackModel() + runner._config = Mock(spec_set=["max_beam_width"], max_beam_width=1) + runner.sample_in_graph_callable = None + runner.keep_greedy_only_outputs = keep + monkeypatch.setattr(decoder_runner_module, "greedy_tokens_usable", lambda request: True) + hidden = torch.ones(2, 4, device="cuda") + tokens = torch.zeros(2, dtype=torch.int32, device="cuda") + graph_outputs = {GREEDY_TOKENS_KEY: tokens, GREEDY_HIDDEN_STATES_KEY: hidden} + + outputs = runner._logits_for_greedy_tokens(_ScheduledGenerationOnly(), graph_outputs) + + if keep: + assert outputs is graph_outputs and "logits" not in outputs + else: + assert torch.equal(outputs["logits"], hidden * 2) + + +def test_an_idle_batch_gets_logits_for_its_dummy_row(): + """An empty batch (an idle attention-DP rank's dummy row) never keeps greedy-only outputs.""" + runner = object.__new__(DecoderRunner) + runner.model = _GreedyFallbackModel() + runner._config = Mock(spec_set=["max_beam_width"], max_beam_width=1) + runner.sample_in_graph_callable = None + hidden = torch.ones(1, 4, device="cuda") + tokens = torch.zeros(1, dtype=torch.int32, device="cuda") + outputs = runner._logits_for_greedy_tokens( + _ScheduledIdle(), {GREEDY_TOKENS_KEY: tokens, GREEDY_HIDDEN_STATES_KEY: hidden} + ) + assert torch.equal(outputs["logits"], hidden * 2) diff --git a/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py b/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py new file mode 100644 index 000000000000..518f5ab9b4ca --- /dev/null +++ b/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DecoderRunner wiring of the step-loop hooks a model declares.""" + +import pytest + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager import ( + MiniMaxM3KVCacheManagerV2, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.runtime_hooks import ( + DEFAULT_RUNTIME_HOOKS, + MsaRuntimePolicy, + RuntimeHooks, +) +from tensorrt_llm.llmapi.llm_args import ( + BaseSparseAttentionConfig, + DeepSeekSparseAttentionConfig, + MiniMaxM3SparseAttentionConfig, +) + +pytestmark = pytest.mark.cpu_only + +_SHORT_IDENTITY = RuntimeHooks(msa=MsaRuntimePolicy(short_sequence_identity=True)) + + +def _runner( + sparse_attention_config: BaseSparseAttentionConfig | None, hooks: RuntimeHooks +) -> DecoderRunner: + # Only the two attributes the threshold and the MSA arming read. + runner = object.__new__(DecoderRunner) + runner.sparse_attention_config = sparse_attention_config + runner.runtime_hooks = hooks + return runner + + +def test_no_threshold_without_a_split_or_a_short_identity_policy(): + assert _runner(None, DEFAULT_RUNTIME_HOOKS)._short_long_seq_len_threshold() is None + assert ( + _runner( + MiniMaxM3SparseAttentionConfig(), DEFAULT_RUNTIME_HOOKS + )._short_long_seq_len_threshold() + is None + ) + + +def test_a_sparse_config_that_asks_for_the_split_sets_the_threshold(): + config = DeepSeekSparseAttentionConfig(index_topk=512, skip_indexer_for_short_seqs=True) + + # The sparse config wins over the hooks. + assert _runner(config, _SHORT_IDENTITY)._short_long_seq_len_threshold() == 512 + + +def test_short_identity_threshold_is_the_topk_span(): + config = MiniMaxM3SparseAttentionConfig(sparse_topk_blocks=16, sparse_block_size=128) + + assert _runner(config, _SHORT_IDENTITY)._short_long_seq_len_threshold() == 16 * 128 + + +@pytest.mark.parametrize( + "config", + [None, DeepSeekSparseAttentionConfig(index_topk=512, skip_indexer_for_short_seqs=False)], + ids=["dense", "dsa_without_split"], +) +def test_short_identity_requires_the_minimax_m3_sparse_config(config): + with pytest.raises(ValueError, match="requires a MiniMax-M3 sparse attention config"): + _runner(config, _SHORT_IDENTITY)._short_long_seq_len_threshold() + + +def test_a_default_msa_policy_arms_nothing(): + runner = _runner(None, DEFAULT_RUNTIME_HOOKS) + # A non-MSA metadata and no cache manager are fine when nothing is armed. + runner.attn_metadata = object.__new__(TrtllmAttentionMetadata) + + runner._arm_msa_policy(None) + + +def test_an_msa_policy_requires_the_msa_backend_and_cache(): + runner = _runner(None, RuntimeHooks(msa=MsaRuntimePolicy(reuse_pool_views=True))) + runner.attn_metadata = object.__new__(TrtllmAttentionMetadata) + + with pytest.raises(ValueError, match="requires the MiniMax-M3 MSA sparse attention backend"): + runner._arm_msa_policy(object.__new__(MiniMaxM3KVCacheManagerV2)) + + +def test_an_msa_policy_reaches_the_metadata_and_the_cache(): + policy = MsaRuntimePolicy( + short_sequence_identity=True, + skip_unchanged_block_offsets=True, + reuse_pool_views=False, + ) + runner = _runner(None, RuntimeHooks(msa=policy)) + runner.attn_metadata = object.__new__(MiniMaxM3MsaSparseAttentionMetadata) + manager = object.__new__(MiniMaxM3KVCacheManagerV2) + + runner._arm_msa_policy(manager) + + assert runner.attn_metadata.msa_policy is policy + assert manager.skip_unchanged_block_offsets is True + assert manager.reuse_pool_views is False diff --git a/tests/unittest/_torch/executor/test_cuda_graph_runner.py b/tests/unittest/_torch/executor/test_cuda_graph_runner.py index 163e1c5449de..b0934702e4f9 100644 --- a/tests/unittest/_torch/executor/test_cuda_graph_runner.py +++ b/tests/unittest/_torch/executor/test_cuda_graph_runner.py @@ -6,10 +6,19 @@ import pytest import torch +from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDAGraphRunner, + CUDAGraphRunnerConfig, EncoderCUDAGraphRunner, EncoderCUDAGraphRunnerConfig, ) +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.bindings import SamplingConfig +from tensorrt_llm.llmapi.llm_args import BaseSparseAttentionConfig, DeepSeekSparseAttentionConfig +from tensorrt_llm.mapping import Mapping pytestmark = pytest.mark.cpu_only @@ -158,3 +167,110 @@ def test_encoder_decoder_fixed_slots_restore_source_order() -> None: restored_output = runner.restore_encoder_decoder_output(key, fixed_slot_output, inputs) expected_output = torch.cat((fixed_slot_output[511:512], fixed_slot_output[:400])) torch.testing.assert_close(restored_output, expected_output) + + +def _decoder_runner( + *, + short_long_seq_len_threshold: int | None = None, + sparse_attention_config: BaseSparseAttentionConfig | None = None, + private_metadata_buffers: bool = False, +) -> CUDAGraphRunner: + """A generation-graph runner with capture disabled, so it allocates no device buffers.""" + config = CUDAGraphRunnerConfig( + use_cuda_graph=False, + cuda_graph_padding_enabled=False, + cuda_graph_batch_sizes=[1, 2], + max_cuda_graph_batch_size=2, + batch_size=2, + max_beam_width=1, + max_num_tokens=2, + use_mrope=False, + spec_config=None, + cuda_graph_mem_pool=None, + enable_attention_dp=False, + original_max_draft_len=0, + original_max_total_draft_tokens=0, + is_encoder_decoder=False, + mapping=Mapping(), + dist=None, + kv_cache_manager_key=ResourceManagerType.KV_CACHE_MANAGER, + sparse_attention_config=sparse_attention_config, + short_long_seq_len_threshold=short_long_seq_len_threshold, + private_metadata_buffers=private_metadata_buffers, + ) + return CUDAGraphRunner(config) + + +def _dsa_config(skip_indexer_for_short_seqs: bool) -> DeepSeekSparseAttentionConfig: + return DeepSeekSparseAttentionConfig( + index_topk=512, skip_indexer_for_short_seqs=skip_indexer_for_short_seqs + ) + + +def test_short_long_threshold_from_the_runner_config_wins() -> None: + runner = _decoder_runner( + short_long_seq_len_threshold=2_048, sparse_attention_config=_dsa_config(True) + ) + + assert runner.short_long_seq_len_threshold == 2_048 + + +def test_short_long_threshold_is_derived_from_a_sparse_config_that_asks_for_the_split() -> None: + runner = _decoder_runner(sparse_attention_config=_dsa_config(True)) + + assert runner.short_long_seq_len_threshold == 512 + + +@pytest.mark.parametrize( + "sparse_attention_config", + [None, _dsa_config(False)], + ids=["dense", "dsa_without_split"], +) +def test_short_long_threshold_is_none_for_a_single_graph_family(sparse_attention_config) -> None: + runner = _decoder_runner(sparse_attention_config=sparse_attention_config) + + assert runner.short_long_seq_len_threshold is None + + +def test_graph_metadata_uses_the_process_buffer_pool_by_default() -> None: + assert _decoder_runner().cuda_graph_meta_buffers is get_memory_buffers() + + +def test_private_metadata_buffers_give_each_runner_its_own_pool() -> None: + first = _decoder_runner(private_metadata_buffers=True) + second = _decoder_runner(private_metadata_buffers=True) + + assert first.cuda_graph_meta_buffers is not get_memory_buffers() + assert first.cuda_graph_meta_buffers is not second.cuda_graph_meta_buffers + + +def _generation_batch(*num_tokens: int) -> ScheduledRequests: + batch = ScheduledRequests() + batch.generation_requests = [ + LlmRequest( + request_id=request_id, + max_new_tokens=4, + input_tokens=list(range(tokens)), + sampling_config=SamplingConfig(), + seq_slot=request_id, + is_streaming=False, + ) + for request_id, tokens in enumerate(num_tokens) + ] + return batch + + +@pytest.mark.parametrize( + ("num_tokens", "short"), + [((100,), True), ((100, 2_048), True), ((100, 2_049), False), ((4_000,), False)], +) +def test_seq_len_mode_splits_at_the_configured_threshold(num_tokens, short) -> None: + runner = _decoder_runner(short_long_seq_len_threshold=2_048) + + assert runner._get_seq_len_mode(_generation_batch(*num_tokens)) is short + + +def test_seq_len_mode_is_long_without_a_threshold() -> None: + runner = _decoder_runner() + + assert runner._get_seq_len_mode(_generation_batch(1)) is False diff --git a/tests/unittest/_torch/executor/test_h2d_cache.py b/tests/unittest/_torch/executor/test_h2d_cache.py new file mode 100644 index 000000000000..d70f58315ff2 --- /dev/null +++ b/tests/unittest/_torch/executor/test_h2d_cache.py @@ -0,0 +1,316 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Skipping host-to-device metadata copies whose values the destination already holds.""" + +import threading + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor import h2d_cache +from tensorrt_llm._torch.pyexecutor.h2d_cache import ( + H2DSkipCache, + copy_h2d_if_changed, + fill_tail_if_changed, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture +def cache(): + """A fresh cache, active for the test body.""" + skip_cache = H2DSkipCache() + with h2d_cache.active(skip_cache): + yield skip_cache + torch.cuda.synchronize() + + +def _host(values, dtype=torch.int32): + return torch.tensor(values, dtype=dtype).pin_memory() + + +def _device_zeros(n, dtype=torch.int32): + return torch.zeros(n, dtype=dtype, device="cuda") + + +def test_without_an_active_cache_every_copy_is_issued_like_copy_(): + assert not h2d_cache.enabled() + dst = _device_zeros(4) + reference = _device_zeros(4) + src = _host([3, 1, 4, 1]) + + for _ in range(3): + assert copy_h2d_if_changed(dst, src) is True + reference.copy_(src, non_blocking=True) + torch.cuda.synchronize() + assert torch.equal(dst, reference) + + +def test_active_scopes_the_cache_to_the_with_block(): + skip_cache = H2DSkipCache() + dst = _device_zeros(2) + values = _host([1, 2]) + + with h2d_cache.active(skip_cache): + assert h2d_cache.enabled() + assert copy_h2d_if_changed(dst, values) is True + assert copy_h2d_if_changed(dst, values) is False + assert not h2d_cache.enabled() + assert copy_h2d_if_changed(dst, values) is True + # Re-activating the same cache finds its shadow again. + with h2d_cache.active(skip_cache): + assert copy_h2d_if_changed(dst, values) is False + + +def test_active_none_disables_skipping_inside_an_outer_cache(cache): + dst = _device_zeros(2) + values = _host([1, 2]) + assert copy_h2d_if_changed(dst, values) is True + + with h2d_cache.active(None): + assert not h2d_cache.enabled() + assert copy_h2d_if_changed(dst, values) is True + assert copy_h2d_if_changed(dst, values) is False + + +def test_two_caches_do_not_share_shadows(cache): + dst = _device_zeros(3) + values = _host([7, 8, 9]) + other = H2DSkipCache() + + assert copy_h2d_if_changed(dst, values) is True + with h2d_cache.active(other): + # The other cache has never seen dst. + assert copy_h2d_if_changed(dst, values) is True + assert copy_h2d_if_changed(dst, values) is False + assert copy_h2d_if_changed(dst, values) is False + + +def test_a_cache_active_on_one_thread_is_invisible_to_another(cache): + dst = _device_zeros(3) + values = _host([4, 5, 6]) + assert copy_h2d_if_changed(dst, values) is True + seen = {} + + def other_thread(): + seen["enabled"] = h2d_cache.enabled() + seen["copied"] = [copy_h2d_if_changed(dst, values) for _ in range(2)] + torch.cuda.synchronize() + + thread = threading.Thread(target=other_thread) + thread.start() + thread.join() + + assert seen == {"enabled": False, "copied": [True, True]} + assert copy_h2d_if_changed(dst, values) is False + + +def test_threads_with_their_own_caches_keep_them_apart(): + dst = [_device_zeros(2) for _ in range(2)] + values = _host([1, 1]) + caches = [H2DSkipCache(), H2DSkipCache()] + barrier = threading.Barrier(2) + results = {} + + def worker(index): + with h2d_cache.active(caches[index]): + first = copy_h2d_if_changed(dst[index], values) + barrier.wait() + second = copy_h2d_if_changed(dst[index], values) + # The other thread's destination was never shadowed by this cache. + other = copy_h2d_if_changed(dst[1 - index], values) + torch.cuda.synchronize() + results[index] = (first, second, other) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert results == {0: (True, False, True), 1: (True, False, True)} + + +def test_a_repeated_identical_copy_is_skipped(cache): + dst = _device_zeros(4) + + assert copy_h2d_if_changed(dst, _host([5, 6, 7, 8])) is True + assert copy_h2d_if_changed(dst, _host([5, 6, 7, 8])) is False + torch.cuda.synchronize() + assert dst.tolist() == [5, 6, 7, 8] + + +def test_clear_makes_the_next_identical_copy_issue(cache): + """The runner clears the cache after a discarded step, whose copies never landed.""" + dst = _device_zeros(3) + values = _host([1, 2, 3]) + assert copy_h2d_if_changed(dst, values) is True + + cache.clear() + + assert copy_h2d_if_changed(dst, values) is True + assert copy_h2d_if_changed(dst, values) is False + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3] + + +def test_changed_values_are_copied(cache): + dst = _device_zeros(4) + + assert copy_h2d_if_changed(dst, _host([1, 2, 3, 4])) is True + assert copy_h2d_if_changed(dst, _host([1, 2, 3, 5])) is True + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3, 5] + assert copy_h2d_if_changed(dst, _host([1, 2, 3, 5])) is False + + +def test_the_shadow_is_a_snapshot_of_the_source(cache): + """Mutating the host source after the copy must not make a later copy look unchanged.""" + dst = _device_zeros(3) + src = _host([1, 2, 3]) + + assert copy_h2d_if_changed(dst, src) is True + torch.cuda.synchronize() + src[0] = 9 + assert copy_h2d_if_changed(dst, src) is True + torch.cuda.synchronize() + assert dst.tolist() == [9, 2, 3] + + +def test_a_source_dtype_change_copies(cache): + dst = _device_zeros(3) + + assert copy_h2d_if_changed(dst, _host([1, 2, 3], torch.int32)) is True + assert copy_h2d_if_changed(dst, _host([1, 2, 3], torch.int64)) is True + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3] + + +def test_a_copy_through_another_view_replaces_the_shadow(cache): + buf = _device_zeros(8) + head, tail = buf[:4], buf[4:] + values = _host([1, 2, 3, 4]) + + assert copy_h2d_if_changed(head, values) is True + assert copy_h2d_if_changed(tail, _host([5, 6, 7, 8])) is True + # The shadow now describes the tail, so the head copy must not be skipped. + assert copy_h2d_if_changed(head, values) is True + assert copy_h2d_if_changed(head, values) is False + torch.cuda.synchronize() + assert buf.tolist() == [1, 2, 3, 4, 5, 6, 7, 8] + + +def test_an_overlapping_view_never_skips_on_stale_state(cache): + buf = _device_zeros(6) + values = _host([1, 1, 1, 1]) + + assert copy_h2d_if_changed(buf[:4], values) is True + assert copy_h2d_if_changed(buf[2:6], _host([2, 2, 2, 2])) is True + # buf[2:4] was overwritten through the other view: same values, same view, must copy. + assert copy_h2d_if_changed(buf[:4], values) is True + torch.cuda.synchronize() + assert buf.tolist() == [1, 1, 1, 1, 2, 2] + + +def test_a_view_with_another_shape_of_the_same_bytes_copies(cache): + buf = _device_zeros(4) + + assert copy_h2d_if_changed(buf, _host([1, 2, 3, 4])) is True + assert copy_h2d_if_changed(buf.view(2, 2), _host([[1, 2], [3, 4]])) is True + assert copy_h2d_if_changed(buf.view(2, 2), _host([[1, 2], [3, 4]])) is False + + +def test_a_cuda_source_always_copies_and_drops_the_shadow(cache): + dst = _device_zeros(4) + values = _host([1, 2, 3, 4]) + device_src = torch.arange(4, dtype=torch.int32, device="cuda") + + assert copy_h2d_if_changed(dst, values) is True + assert copy_h2d_if_changed(dst, device_src) is True + assert copy_h2d_if_changed(dst, device_src) is True + # The device-side write replaced what the shadow described. + assert copy_h2d_if_changed(dst, values) is True + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3, 4] + + +def test_a_host_destination_always_copies(cache): + dst = torch.zeros(3, dtype=torch.int32) + + assert copy_h2d_if_changed(dst, _host([1, 2, 3])) is True + assert copy_h2d_if_changed(dst, _host([1, 2, 3])) is True + assert dst.tolist() == [1, 2, 3] + + +def test_capture_always_copies_and_drops_the_shadow(cache): + dst = _device_zeros(4) + values = _host([4, 3, 2, 1]) + assert copy_h2d_if_changed(dst, values) is True + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + results = [] + with torch.cuda.graph(graph): + results.append(copy_h2d_if_changed(dst, values)) + assert results == [True] + + dst.zero_() + graph.replay() + torch.cuda.synchronize() + assert dst.tolist() == [4, 3, 2, 1] + # The captured copy reruns on every replay, so the shadow cannot vouch for the buffer. + assert copy_h2d_if_changed(dst, values) is True + + +def test_without_an_active_cache_every_tail_fill_runs(): + buf = _device_zeros(6) + fill_tail_if_changed(buf, 2, -1) + buf[3] = 5 + fill_tail_if_changed(buf, 2, -1) + fill_tail_if_changed(buf, 6, -1) + assert buf.tolist() == [0, 0, -1, -1, -1, -1] + + +def test_a_repeated_tail_fill_is_skipped_and_a_moved_start_refills(cache): + buf = _device_zeros(6) + fill_tail_if_changed(buf, 2, -1) + # Writes nothing else is allowed to make; they show which fills ran. + buf[3] = 5 + fill_tail_if_changed(buf, 2, -1) + assert buf[3].item() == 5 + fill_tail_if_changed(buf, 4, -1) + buf[3] = 7 + fill_tail_if_changed(buf, 2, -1) + assert buf.tolist() == [0, 0, -1, -1, -1, -1] + + +def test_a_whole_buffer_write_drops_the_tail_record(cache): + buf = _device_zeros(6) + fill_tail_if_changed(buf, 2, -1) + buf.copy_(torch.arange(6, dtype=torch.int32, device="cuda")) + fill_tail_if_changed(buf, 6, -1) + fill_tail_if_changed(buf, 2, -1) + assert buf.tolist() == [0, 1, -1, -1, -1, -1] + + +def test_clear_drops_the_tail_records(cache): + buf = _device_zeros(6) + fill_tail_if_changed(buf, 2, -1) + buf[4] = 9 + cache.clear() + fill_tail_if_changed(buf, 2, -1) + assert buf[2:].eq(-1).all() + + +def test_a_tail_record_keeps_its_storage_so_a_new_buffer_cannot_inherit_it(cache): + buf = _device_zeros(1 << 20) + fill_tail_if_changed(buf, 3, -1) + address = buf.data_ptr() + del buf + # Without the record holding the storage, the allocator hands this same block back, and the fill below would be + # skipped on the freed buffer's record. + fresh = _device_zeros(1 << 20) + assert fresh.data_ptr() != address + fill_tail_if_changed(fresh, 3, -1) + assert fresh[3:].eq(-1).all() diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 7c9398a85892..86805f3204a7 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -102,6 +102,7 @@ def test_forward_step_carries_context_logits_request_to_runner(request_flags, ex iter_counter=0, sampler=Mock(get_cache_indirection=Mock(return_value=cache_indirection)), execution_stream=Mock(), + _loop_on_execution_stream=False, _iter_adp_dummy_ctx_tokens=0, _iter_adp_dummy_gen_tokens=0, _compute_adp_dummy_tokens=PyExecutor._compute_adp_dummy_tokens, diff --git a/tests/unittest/_torch/executor/test_runtime_hooks.py b/tests/unittest/_torch/executor/test_runtime_hooks.py new file mode 100644 index 000000000000..2dadaa957e8e --- /dev/null +++ b/tests/unittest/_torch/executor/test_runtime_hooks.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""How the engine reads the step-loop hooks a model declares (``runtime_hooks_of``).""" + +import dataclasses + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.runtime_hooks import ( + DEFAULT_RUNTIME_HOOKS, + MsaRuntimePolicy, + RuntimeHooks, + runtime_hooks_of, +) + +pytestmark = pytest.mark.cpu_only + + +class _PlainModel(torch.nn.Module): + """A model that declares no hooks, like every model outside the opt-in.""" + + +class _ClassHooksModel(torch.nn.Module): + runtime_hooks = RuntimeHooks(pack_step_inputs=True, skip_unchanged_h2d=True) + + +def test_model_without_hooks_gets_the_all_off_default(): + hooks = runtime_hooks_of(_PlainModel()) + + assert hooks is DEFAULT_RUNTIME_HOOKS + assert hooks == RuntimeHooks() + assert not hooks.pack_step_inputs + assert not hooks.skip_unchanged_h2d + assert not hooks.loop_on_execution_stream + assert not hooks.merge_tuned_tactics_across_ranks + assert not hooks.skip_greedy_rng_upload + assert hooks.spec_step_fold is None + assert hooks.msa == MsaRuntimePolicy() + assert not hooks.msa.short_sequence_identity + assert hooks.msa.kv_update is None + assert not hooks.msa.skip_unchanged_block_offsets + assert not hooks.msa.reuse_pool_views + + +def test_hooks_declared_as_none_mean_the_default(): + model = _PlainModel() + model.runtime_hooks = None + + assert runtime_hooks_of(model) is DEFAULT_RUNTIME_HOOKS + + +def test_class_attribute_hooks_are_returned_as_declared(): + model = _ClassHooksModel() + + assert runtime_hooks_of(model) is _ClassHooksModel.runtime_hooks + + +def test_instance_attribute_hooks_are_returned_as_declared(): + hooks = RuntimeHooks( + loop_on_execution_stream=True, + msa=MsaRuntimePolicy(short_sequence_identity=True, reuse_pool_views=True), + ) + model = _PlainModel() + model.runtime_hooks = hooks + + assert runtime_hooks_of(model) is hooks + + +def test_instance_attribute_overrides_class_attribute(): + hooks = RuntimeHooks(skip_greedy_rng_upload=True) + model = _ClassHooksModel() + model.runtime_hooks = hooks + + assert runtime_hooks_of(model) is hooks + + +@pytest.mark.parametrize( + "declared", + [True, {"pack_step_inputs": True}, MsaRuntimePolicy(), "RuntimeHooks()"], + ids=["bool", "dict", "msa_policy", "str"], +) +def test_hooks_of_the_wrong_type_raise(declared): + model = _PlainModel() + model.runtime_hooks = declared + + with pytest.raises(TypeError, match=r"_PlainModel\.runtime_hooks must be a RuntimeHooks"): + runtime_hooks_of(model) + + +def test_hooks_are_immutable_once_declared(): + hooks = RuntimeHooks() + + with pytest.raises(dataclasses.FrozenInstanceError): + hooks.pack_step_inputs = True + with pytest.raises(dataclasses.FrozenInstanceError): + hooks.msa.reuse_pool_views = True + + +@pytest.mark.parametrize( + "hooks, skips", + [ + (RuntimeHooks(), False), + (RuntimeHooks(pack_step_inputs=True, skip_greedy_rng_upload=True), False), + (RuntimeHooks(skip_unchanged_h2d=True), True), + (RuntimeHooks(msa=MsaRuntimePolicy(skip_unchanged_block_offsets=True)), True), + ], + ids=["default", "no_skip_of_metadata", "h2d_skip", "msa_block_offset_skip"], +) +def test_every_metadata_upload_skip_asks_for_private_metadata_buffers(hooks, skips): + assert hooks.skips_unchanged_metadata_uploads is skips diff --git a/tests/unittest/_torch/executor/test_step_inputs.py b/tests/unittest/_torch/executor/test_step_inputs.py new file mode 100644 index 000000000000..f7f3859a5d05 --- /dev/null +++ b/tests/unittest/_torch/executor/test_step_inputs.py @@ -0,0 +1,604 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Packing a decode step's small input copies into one upload and one scatter launch.""" + +import threading + +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 + + +def _scatter_op_registered() -> bool: + return torch.cuda.is_available() and hasattr(torch.ops.trtllm, "step_inputs_scatter") + + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.skipif( + not _scatter_op_registered(), + reason="torch.ops.trtllm.step_inputs_scatter is not registered (built from " + "cpp/tensorrt_llm/thop/stepInputsScatterOp.cpp)", + ), +] + +# About half a second of GPU time: long enough that work queued behind it is still pending when +# the host checks, short enough for a unit test. +_SLEEP_CYCLES = 1_000_000_000 + + +def _device(): + return torch.device("cuda", torch.cuda.current_device()) + + +@pytest.fixture +def pack(): + """A fresh pack, active for the test body.""" + step_pack = StepInputPack(_device()) + with step_inputs.active(step_pack): + yield step_pack + step_inputs.flush() + torch.cuda.synchronize() + + +def _host(values, dtype=torch.int32): + return torch.tensor(values, dtype=dtype).pin_memory() + + +def _zeros(*shape, dtype=torch.int32): + return torch.zeros(*shape, dtype=dtype, device="cuda") + + +# --------------------------------------------------------------------------- +# No active pack +# --------------------------------------------------------------------------- + + +def test_without_an_active_pack_every_call_is_a_plain_copy(): + dst = _zeros(4) + d2d_dst = _zeros(4) + src_dev = torch.arange(4, dtype=torch.int32, device="cuda") + + step_inputs.begin() + assert not step_inputs.is_open() + assert step_inputs.copy_(dst, _host([1, 2, 3, 4])) is dst + step_inputs.copy_(d2d_dst, src_dev) + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3, 4] + assert d2d_dst.tolist() == [0, 1, 2, 3] + step_inputs.flush() + step_inputs.discard() + + +def test_an_inactive_pack_is_not_used(): + idle = StepInputPack(_device()) + dst = _zeros(2) + + step_inputs.copy_(dst, _host([3, 4])) + torch.cuda.synchronize() + + assert dst.tolist() == [3, 4] + assert not idle.open and idle.staged == [] + + +# --------------------------------------------------------------------------- +# Deferral and ordering +# --------------------------------------------------------------------------- + + +def test_packed_host_copy_lands_only_at_flush(pack): + dst = _zeros(4) + src = _host([1, 2, 3, 4]) + + step_inputs.begin() + assert step_inputs.is_open() and pack.open + assert step_inputs.copy_(dst, src) is dst + # The pack snapshots the host bytes at the copy site. + src.fill_(7) + torch.cuda.synchronize() + assert dst.tolist() == [0, 0, 0, 0] + + step_inputs.flush() + assert not step_inputs.is_open() + assert pack.staged == [] + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3, 4] + + +def test_packed_host_copy_into_a_row_strided_2d_slice(pack): + table = torch.full((4, 6), -1, dtype=torch.int32, device="cuda") + dst = table[1:3, 2:5] + assert not dst.is_contiguous() and dst.stride(1) == 1 + expected = table.clone() + expected[1:3, 2:5] = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32) + + step_inputs.begin() + step_inputs.copy_(dst, _host([[1, 2, 3], [4, 5, 6]])) + torch.cuda.synchronize() + assert torch.equal(table, torch.full_like(table, -1)) + + step_inputs.flush() + torch.cuda.synchronize() + assert torch.equal(table, expected) + + +def test_packed_host_copy_converts_dtype_on_the_host(pack): + dst = _zeros(3) + + step_inputs.begin() + step_inputs.copy_(dst, _host([7, -8, 9], dtype=torch.int64)) + step_inputs.flush() + torch.cuda.synchronize() + + assert dst.tolist() == [7, -8, 9] + + +def test_packed_device_copy_lands_only_at_flush(pack): + dst = _zeros(5) + src = torch.arange(5, dtype=torch.int32, device="cuda") + + step_inputs.begin() + step_inputs.copy_(dst, src) + torch.cuda.synchronize() + assert dst.tolist() == [0] * 5 + + step_inputs.flush() + torch.cuda.synchronize() + assert dst.tolist() == [0, 1, 2, 3, 4] + + +def test_device_copy_reads_what_a_host_copy_of_the_same_pack_wrote(pack): + """The graph runner copies the engine's freshly uploaded buffers into its static inputs.""" + staged = _zeros(4) + static = _zeros(1, 8) + + step_inputs.begin() + step_inputs.copy_(staged, _host([10, 11, 12, 13])) + step_inputs.copy_(static[:, :4], staged.unsqueeze(0)) + step_inputs.flush() + torch.cuda.synchronize() + + assert static.tolist() == [[10, 11, 12, 13, 0, 0, 0, 0]] + + +# The scatter's CTA copies descriptor i with thread t on bytes t, t + 256, ... of it. When the +# second copy starts TAIL bytes before the end of the first (TAIL not a multiple of the warp +# size), its first bytes are the ones the first copy writes last, from other warps: without a +# barrier between descriptors, a warp that finished the first copy early writes (or reads) them +# before a slower warp has written its part of the first copy. +_TAILS = [239, 1000, 4073] + + +@pytest.mark.parametrize("tail", _TAILS) +@pytest.mark.parametrize("kind", ["host", "device"]) +def test_overlapping_destinations_land_in_descriptor_order(pack, kind, tail): + """The later copy wins on the overlap, as it would eagerly.""" + # Host copies share the 64 KiB staging buffer; device copies are bounded only by memory, and a + # long first copy lets the warps drift apart. + n = 24 * 1024 if kind == "host" else 1 << 20 + m = 8 * 1024 + first = torch.full((n,), 1, dtype=torch.uint8) + second = torch.full((m,), 2, dtype=torch.uint8) + if kind == "host": + first, second = first.pin_memory(), second.pin_memory() + else: + first, second = first.cuda(), second.cuda() + start = n - tail + expected = torch.ones(start + m, dtype=torch.uint8) + expected[start:] = 2 + + mismatches = [] + for trial in range(50): + buf = torch.zeros(start + m, dtype=torch.uint8, device="cuda") + step_inputs.begin() + step_inputs.copy_(buf[:n], first) + step_inputs.copy_(buf[start:], second) + step_inputs.flush() + torch.cuda.synchronize() + wrong = int((buf.cpu() != expected).sum()) + if wrong: + mismatches.append((trial, wrong)) + + assert not mismatches, f"(trial, overlap bytes where the earlier copy won): {mismatches}" + + +@pytest.mark.parametrize("tail", _TAILS) +def test_a_device_copy_reads_the_destination_of_an_earlier_device_copy(pack, tail): + """A chain of device copies: the second reads the bytes the first writes last.""" + n = 1 << 20 + source = torch.randint(1, 256, (n,), dtype=torch.uint8, device="cuda") + expected = source[n - tail :].clone() + + mismatches = [] + for trial in range(50): + middle = torch.zeros(n, dtype=torch.uint8, device="cuda") + out = torch.zeros(tail, dtype=torch.uint8, device="cuda") + step_inputs.begin() + step_inputs.copy_(middle, source) + step_inputs.copy_(out, middle[n - tail :]) + step_inputs.flush() + torch.cuda.synchronize() + wrong = int((out != expected).sum()) + if wrong: + mismatches.append((trial, wrong)) + + assert not mismatches, ( + f"(trial, stale bytes read before the first copy wrote them): {mismatches}" + ) + + +# --------------------------------------------------------------------------- +# Staged tensors stay alive +# --------------------------------------------------------------------------- + + +def test_a_temporary_device_source_outlives_its_copy(pack): + """The pack holds the source until the scatter is queued, so its memory is not reused.""" + numel = 1 << 16 + expected = torch.arange(numel, dtype=torch.int32) * 3 + 1 + dst = _zeros(numel) + + step_inputs.begin() + # The only reference to this source is the one the pack keeps. + step_inputs.copy_(dst, torch.arange(numel, dtype=torch.int32, device="cuda") * 3 + 1) + assert len(pack.staged) == 2 + # Churn the allocator with same-sized blocks: a freed source block would be handed out here. + garbage = [torch.full((numel,), -7, dtype=torch.int32, device="cuda") for _ in range(16)] + torch.cuda.synchronize() + step_inputs.flush() + del garbage + # Blocks allocated after the flush are ordered behind the scatter on this stream. + after = [torch.full((numel,), -9, dtype=torch.int32, device="cuda") for _ in range(16)] + torch.cuda.synchronize() + + assert torch.equal(dst.cpu(), expected) + del after + + +def test_a_temporary_destination_view_outlives_its_copy(pack): + owner = _zeros(8) + + step_inputs.begin() + step_inputs.copy_(owner.view(2, 4)[1], _host([5, 6, 7, 8])) + assert len(pack.staged) == 1 + step_inputs.flush() + torch.cuda.synchronize() + + assert owner.tolist() == [0, 0, 0, 0, 5, 6, 7, 8] + + +# --------------------------------------------------------------------------- +# Fallbacks +# --------------------------------------------------------------------------- + + +def test_an_unpackable_copy_flushes_first_and_keeps_order(pack): + buf = _zeros(8) + other = _zeros(4) + later = _zeros(2) + + step_inputs.begin() + step_inputs.copy_(buf, _host([1] * 8)) + step_inputs.copy_(other, _host([3, 3, 3, 3])) + # A strided 1-D destination cannot be packed: the pack lands first, then this copy. + step_inputs.copy_(buf[::2], _host([2, 2, 2, 2])) + assert step_inputs.is_open() + torch.cuda.synchronize() + assert buf.tolist() == [2, 1, 2, 1, 2, 1, 2, 1] + assert other.tolist() == [3, 3, 3, 3] + + # The pack reopened: later copies are deferred again. + step_inputs.copy_(later, _host([5, 6])) + torch.cuda.synchronize() + assert later.tolist() == [0, 0] + step_inputs.flush() + torch.cuda.synchronize() + assert later.tolist() == [5, 6] + + +def test_a_device_copy_with_a_dtype_change_is_not_packed(pack): + dst = _zeros(3, dtype=torch.int64) + src = torch.tensor([4, 5, 6], dtype=torch.int32, device="cuda") + + step_inputs.begin() + step_inputs.copy_(dst, src) + torch.cuda.synchronize() + + assert dst.tolist() == [4, 5, 6] + assert step_inputs.is_open() + + +def test_a_copy_larger_than_the_staging_buffer_falls_back_in_order(pack): + numel = StepInputPack.CAPACITY // 4 + big = _zeros(numel) + big_values = torch.arange(numel, dtype=torch.int32).pin_memory() + + step_inputs.begin() + step_inputs.copy_(big[:4], _host([-1, -1, -1, -1])) + step_inputs.copy_(big, big_values) + torch.cuda.synchronize() + + # The staged prefix landed before the full copy, which therefore wins. + assert torch.equal(big.cpu(), big_values) + assert step_inputs.is_open() + + +def test_copies_that_together_overflow_the_staging_buffer_all_land(pack): + # Each copy is a quarter of the staging buffer, so the fourth one does not fit. + chunk = StepInputPack.CAPACITY // 16 + dsts = [_zeros(chunk) for _ in range(6)] + values = [torch.full((chunk,), i + 1, dtype=torch.int32).pin_memory() for i in range(6)] + + step_inputs.begin() + for dst, value in zip(dsts, values): + step_inputs.copy_(dst, value) + step_inputs.flush() + torch.cuda.synchronize() + + for i, dst in enumerate(dsts): + assert torch.equal(dst.cpu(), values[i]), f"copy {i} did not land" + + +def test_more_copies_than_descriptors_all_land(pack): + count = StepInputPack.MAX_DESC + 40 + dst = _zeros(count) + + step_inputs.begin() + for i in range(count): + step_inputs.copy_(dst[i : i + 1], _host([i + 1])) + step_inputs.flush() + torch.cuda.synchronize() + + assert dst.tolist() == list(range(1, count + 1)) + + +def test_a_row_strided_copy_that_needs_more_descriptors_than_remain_falls_back(pack): + rows = 8 + table = torch.full((rows, 4), -1, dtype=torch.int32, device="cuda") + filler = _zeros(StepInputPack.MAX_DESC - 2) + rows_values = torch.arange(rows * 3, dtype=torch.int32).reshape(rows, 3) + + step_inputs.begin() + for i in range(filler.numel()): + step_inputs.copy_(filler[i : i + 1], _host([i + 1])) + # Two descriptors remain; the slice needs one per row, so it flushes and copies directly. + step_inputs.copy_(table[:, :3], rows_values.pin_memory()) + step_inputs.flush() + torch.cuda.synchronize() + + assert filler.tolist() == list(range(1, filler.numel() + 1)) + assert torch.equal(table[:, :3].cpu(), rows_values) + assert table[:, 3].tolist() == [-1] * rows + + +def test_mixed_device_copies_past_the_descriptor_limit_all_land(pack): + count = StepInputPack.MAX_DESC + 10 + src = torch.arange(1, count + 1, dtype=torch.int32, device="cuda") + dst = _zeros(count) + + step_inputs.begin() + for i in range(count): + step_inputs.copy_(dst[i : i + 1], src[i : i + 1]) + step_inputs.flush() + torch.cuda.synchronize() + + assert torch.equal(dst, src) + + +def test_begin_on_an_open_pack_issues_it_first(pack): + dst = _zeros(2) + + step_inputs.begin() + step_inputs.copy_(dst, _host([8, 9])) + step_inputs.begin() + torch.cuda.synchronize() + + assert dst.tolist() == [8, 9] + + +def test_reusing_a_pinned_staging_buffer_waits_for_its_upload(pack): + """With every staging buffer in the ring used once, the next pack reuses the first one. + + Its upload is queued behind a long kernel, so rewriting the buffer without waiting would + upload the newer bytes in place of the older ones. + """ + ring = StepInputPack.RING + dsts = [_zeros(4) for _ in range(ring + 1)] + values = [[i * 10 + j for j in range(4)] for i in range(ring + 1)] + + step_inputs.begin() + torch.cuda._sleep(_SLEEP_CYCLES) + step_inputs.copy_(dsts[0], _host(values[0])) + step_inputs.flush() + first_upload = pack.events[0] + for i in range(1, ring): + step_inputs.begin() + step_inputs.copy_(dsts[i], _host(values[i])) + step_inputs.flush() + assert not first_upload.query(), "the queued kernel finished too early for this check" + + step_inputs.begin() + assert first_upload.query() + step_inputs.copy_(dsts[ring], _host(values[ring])) + step_inputs.flush() + torch.cuda.synchronize() + + for i, dst in enumerate(dsts): + assert dst.tolist() == values[i], f"pack {i} uploaded the wrong bytes" + + +def test_a_copy_while_capturing_is_not_deferred(pack): + dst = _zeros(3) + src = torch.tensor([1, 2, 3], dtype=torch.int32, device="cuda") + + step_inputs.begin() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step_inputs.copy_(dst, src) + step_inputs.flush() + torch.cuda.synchronize() + assert dst.tolist() == [0, 0, 0] + + graph.replay() + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3] + + +# --------------------------------------------------------------------------- +# Discard +# --------------------------------------------------------------------------- + + +def test_discard_drops_the_staged_copies(pack): + dst = _zeros(3) + + step_inputs.begin() + step_inputs.copy_(dst, _host([1, 2, 3])) + step_inputs.copy_(dst, torch.ones(3, dtype=torch.int32, device="cuda")) + step_inputs.discard() + assert not pack.open and pack.staged == [] and pack.h2d == [] and pack.d2d == [] + step_inputs.flush() + torch.cuda.synchronize() + + assert dst.tolist() == [0, 0, 0] + + +def test_after_a_discarded_step_identical_values_are_sent_again(pack): + """The runner's recovery: discard the pack and clear the skip cache.""" + skip_cache = H2DSkipCache() + dst = _zeros(3) + values = _host([4, 5, 6]) + + with h2d_cache.active(skip_cache): + step_inputs.begin() + assert copy_h2d_if_changed(dst, values, pack=True) is True + step_inputs.discard() + torch.cuda.synchronize() + assert dst.tolist() == [0, 0, 0] + # Without clearing, the cache would vouch for values that never landed. + assert copy_h2d_if_changed(dst, values, pack=True) is False + skip_cache.clear() + + step_inputs.begin() + assert copy_h2d_if_changed(dst, values, pack=True) is True + step_inputs.flush() + torch.cuda.synchronize() + + assert dst.tolist() == [4, 5, 6] + + +def test_skip_cache_copies_join_the_active_pack(pack): + skip_cache = H2DSkipCache() + dst = _zeros(3) + values = _host([4, 5, 6]) + + with h2d_cache.active(skip_cache): + step_inputs.begin() + assert copy_h2d_if_changed(dst, values, pack=True) is True + torch.cuda.synchronize() + assert dst.tolist() == [0, 0, 0] + # Already staged with these values: the second copy is skipped. + assert copy_h2d_if_changed(dst, values, pack=True) is False + step_inputs.flush() + torch.cuda.synchronize() + + assert dst.tolist() == [4, 5, 6] + + +# --------------------------------------------------------------------------- +# Isolation between packs +# --------------------------------------------------------------------------- + + +def test_a_nested_pack_does_not_see_the_outer_pack(pack): + outer_dst = _zeros(2) + inner_dst = _zeros(2) + other = StepInputPack(_device()) + + step_inputs.begin() + step_inputs.copy_(outer_dst, _host([1, 2])) + with step_inputs.active(other): + assert not step_inputs.is_open() + step_inputs.copy_(inner_dst, _host([3, 4])) # plain copy: the inner pack is closed + step_inputs.begin() + step_inputs.copy_(inner_dst, _host([5, 6])) + torch.cuda.synchronize() + assert inner_dst.tolist() == [3, 4] + assert outer_dst.tolist() == [0, 0] + step_inputs.flush() + assert step_inputs.is_open() and len(pack.h2d) == 1 + step_inputs.flush() + torch.cuda.synchronize() + + assert inner_dst.tolist() == [5, 6] + assert outer_dst.tolist() == [1, 2] + + +def test_active_none_disables_packing_inside_an_outer_pack(pack): + dst = _zeros(2) + + step_inputs.begin() + with step_inputs.active(None): + assert not step_inputs.is_open() + step_inputs.copy_(dst, _host([7, 8])) + torch.cuda.synchronize() + assert dst.tolist() == [7, 8] + assert pack.h2d == [] + + +def test_an_open_pack_on_one_thread_is_invisible_to_another(pack): + mine = _zeros(2) + theirs = _zeros(2) + seen = {} + + step_inputs.begin() + step_inputs.copy_(mine, _host([1, 1])) + + def other_thread(): + seen["open"] = step_inputs.is_open() + step_inputs.copy_(theirs, _host([2, 2])) + step_inputs.flush() + torch.cuda.synchronize() + seen["theirs"] = theirs.tolist() + seen["mine"] = mine.tolist() + + thread = threading.Thread(target=other_thread) + thread.start() + thread.join() + + assert seen == {"open": False, "theirs": [2, 2], "mine": [0, 0]} + assert step_inputs.is_open() and len(pack.h2d) == 1 + step_inputs.flush() + torch.cuda.synchronize() + assert mine.tolist() == [1, 1] + + +def test_threads_with_their_own_packs_keep_them_apart(): + packs = [StepInputPack(_device()) for _ in range(2)] + dsts = [_zeros(4) for _ in range(2)] + barrier = threading.Barrier(2) + results = {} + + def worker(index): + with step_inputs.active(packs[index]): + step_inputs.begin() + step_inputs.copy_(dsts[index], _host([index + 1] * 4)) + barrier.wait() + # Both packs are open with one staged copy each. + results[index] = (packs[index].open, len(packs[index].h2d)) + barrier.wait() + step_inputs.flush() + torch.cuda.synchronize() + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert results == {0: (True, 1), 1: (True, 1)} + assert dsts[0].tolist() == [1] * 4 + assert dsts[1].tolist() == [2] * 4 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/_m3_fused_attn_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/_m3_fused_attn_step.py new file mode 100644 index 000000000000..85eddd0cddc4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/_m3_fused_attn_step.py @@ -0,0 +1,255 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A short-sequence decode step of one MiniMax-M3 TP4 sparse layer, for the fused attention entries' tests +(``m3_fused_decode_attn``, ``m3_attn_mega``, ``m3_attn_c1``). + +The history is written into the paged NVFP4 pool by the in-tree scatter (``fused_write_layer_caches_nvfp4``, which +quantizes with ``torch.ops.trtllm.fp4_quantize``) at the same ``kv_quant_scale`` the step inserts with. The new +tokens' slots and every row past a request's length hold random bytes, so a kernel that skipped the insert or +read past the length is caught. Pages are shuffled so a kernel that ignored the block table reads the wrong ones. +""" + +from __future__ import annotations + +import math +from types import SimpleNamespace +from typing import Sequence + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_fused_decode_attn import ( # noqa: E501 + attend_cached, + nvfp4_dequantize, + query_rows, + v_scale_rows, +) +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 +PACKED_WIDTH = (NUM_HEADS + 4) * HEAD_DIM +HIDDEN = 6144 +#: The MXFP8 hand-off buffers' rows (the target's MERGE_MXFP8 buffer). +MXFP8_ROWS = 128 +#: NVFP4 second-level scales: block scales stay inside E4M3 for |x| up to 16 (K) and 8 (V, N(0, 1) rows). +K_QUANT = 448.0 * 6.0 / 16.0 +V_QUANT = 448.0 * 6.0 / 8.0 +ROPE_THETA = 5.0e6 +MAX_POSITION = 2304 + + +def rotary_cos_sin(device) -> torch.Tensor: + """fp32 [max_pos, 2, 32]: NeoX cos / sin of the 64 rotary dims at every position.""" + inv_freq = 1.0 / ( + ROPE_THETA ** (torch.arange(0, 64, 2, dtype=torch.float64, device=device) / 64) + ) + ang = ( + torch.arange(MAX_POSITION, dtype=torch.float64, device=device)[:, None] * inv_freq[None, :] + ) + return torch.stack([ang.cos(), ang.sin()], dim=1).float().contiguous() + + +def build_step( + seq_lens: Sequence[int], qlen: int, seed: int, packed: torch.Tensor = None +) -> SimpleNamespace: + """Caches, metadata and weights of one step; ``packed`` (bf16 [T, 2560]) defaults to N(0, 1) rows.""" + dev = torch.device("cuda") + g = torch.Generator(device=dev).manual_seed(seed) + batch = len(seq_lens) + t_total = batch * qlen + max_blocks = max(math.ceil(s / PAGE) for s in seq_lens) + num_pages = batch * (max_blocks + 1) + 3 + # One spare column: the kernel must read the columns the length needs, not the table's width. + block_table = torch.randperm(num_pages, device=dev, generator=g)[: batch * (max_blocks + 1)] + block_table = block_table.to(torch.int32).reshape(batch, max_blocks + 1).contiguous() + kv_data = torch.randint( + 0, 256, (num_pages, 2, 1, PAGE, HEAD_DIM // 2), dtype=torch.uint8, device=dev, generator=g + ) + kv_scale = torch.randint( + 0, 256, (num_pages, 2, 1, PAGE, HEAD_DIM // 16), dtype=torch.uint8, device=dev, generator=g + ) + idx_k = torch.randint( + 0, 256, (num_pages, 1, PAGE, HEAD_DIM), dtype=torch.uint8, device=dev, generator=g + ).view(torch.float8_e4m3fn) + kv_quant_scale = torch.tensor([1.0, K_QUANT, V_QUANT], dtype=torch.float32, device=dev) + + hist_slots = [] + for r, n in enumerate(seq_lens): + for pos in range(n - qlen): + hist_slots.append(int(block_table[r, pos // PAGE]) * PAGE + pos % PAGE) + if hist_slots: + h = len(hist_slots) + k_hist = (torch.randn(h, HEAD_DIM, device=dev, generator=g) * 1.5).to(torch.bfloat16) + v_hist = torch.randn(h, HEAD_DIM, device=dev, generator=g).to(torch.bfloat16) + i_hist = torch.randn(h, HEAD_DIM, device=dev, generator=g).to(torch.bfloat16) + assert fused_write_layer_caches_nvfp4( + kv_data[:, 0], + kv_data[:, 1], + kv_scale[:, 0], + kv_scale[:, 1], + idx_k, + torch.tensor(hist_slots, dtype=torch.int32, device=dev), + k_hist, + v_hist, + i_hist, + kv_quant_scale, + ) + + positions, slots = [], [] + for r, n in enumerate(seq_lens): + for t in range(qlen): + pos = n - qlen + t + positions.append(pos) + slots.append(int(block_table[r, pos // PAGE]) * PAGE + pos % PAGE) + if packed is None: + packed = torch.randn(t_total, PACKED_WIDTH, device=dev, generator=g).to(torch.bfloat16) + + def norm_weight(): + return (torch.randn(HEAD_DIM, device=dev, generator=g) * 0.1).to(torch.bfloat16) + + return SimpleNamespace( + packed=packed, + kv_data=kv_data, + kv_scale=kv_scale, + idx_k=idx_k, + out_cache_loc=torch.tensor(slots, dtype=torch.int32, device=dev), + position_ids=torch.tensor(positions, dtype=torch.int32, device=dev), + block_table=block_table, + seq_lens=torch.tensor(list(seq_lens), dtype=torch.int32, device=dev), + kv_quant_scale=kv_quant_scale, + k_global_scale=torch.tensor([1.0 / K_QUANT], dtype=torch.float32, device=dev), + v_global_scale=torch.tensor([1.0 / V_QUANT], dtype=torch.float32, device=dev), + q_norm_weight=norm_weight(), + k_norm_weight=norm_weight(), + index_k_norm_weight=norm_weight(), + rotary_cos_sin=rotary_cos_sin(dev), + eps=1e-6, + sm_scale=HEAD_DIM**-0.5, + decode_query_len=qlen, + num_tokens=t_total, + ) + + +def attention_args(s: SimpleNamespace) -> tuple: + """The attention arguments every entry takes, in schema order (``kv_data`` .. ``rotary_cos_sin``, the two + scalars and ``decode_query_len``).""" + return ( + s.kv_data, + s.kv_scale, + s.idx_k, + s.out_cache_loc, + s.position_ids, + s.block_table, + s.seq_lens, + s.kv_quant_scale, + s.k_global_scale, + s.v_global_scale, + s.q_norm_weight, + s.k_norm_weight, + s.index_k_norm_weight, + s.rotary_cos_sin, + s.eps, + s.sm_scale, + s.decode_query_len, + ) + + +def swizzle(rows_exp: torch.Tensor, padded_rows: int) -> torch.Tensor: + """[rows, k_blocks] UE8M0 bytes -> the 128x4-swizzled buffer of ``padded_rows`` rows (filler 127).""" + rows, kb = rows_exp.shape + k_tiles = (kb + 3) // 4 + out = torch.full((padded_rows * k_tiles * 4,), 127, dtype=torch.uint8, device=rows_exp.device) + r = torch.arange(rows, device=rows_exp.device)[:, None] + k = torch.arange(kb, device=rows_exp.device)[None, :] + off = ( + (r // 128) * k_tiles * 512 + (k // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + k % 4 + ) + out[off.reshape(-1)] = rows_exp.reshape(-1) + return out + + +def mxfp8_operand( + rows: int, cols: int, exp_lo: int, exp_hi: int, g: torch.Generator, padded_rows: int +): + """A random MXFP8 operand: E4M3 values from N(0, 1) and UE8M0 exponents in [exp_lo, exp_hi], swizzled.""" + dev = torch.device("cuda") + q = torch.randn(rows, cols, device=dev, generator=g).to(torch.float8_e4m3fn) + ex = torch.randint( + exp_lo, exp_hi + 1, (rows, cols // 32), dtype=torch.uint8, device=dev, generator=g + ) + return q, swizzle(ex, padded_rows) + + +def new_slot_mask(s: SimpleNamespace, num_pages: int) -> torch.Tensor: + """bool [pages, 128]: the rows the step inserts.""" + mask = torch.zeros(num_pages, PAGE, dtype=torch.bool, device=s.kv_data.device) + for slot in s.out_cache_loc.tolist(): + mask[slot // PAGE, slot % PAGE] = True + return mask + + +# Band of the attention output, in bf16 ulps (2^-8) of each [128] head row's largest magnitude, and of the relative +# RMS, against the float64 decode over the caches the op wrote. Both sides read the same operands (E4M3 q, E2M1 x +# E4M3 K/V, all exact in fp16), so the gap is the bf16 rounding on both sides (half an ulp each), the fp16 +# probabilities before PV (2^-11 relative), fp32 accumulation and approximate exp2/log2 in the split-K partials and +# their LSE merge. Measured over 46 steps: at most 2.0 ulps per element, 0.2 ulp RMS. +ATTN_ELEMENT_ULP = 3.0 +ATTN_RMS_ULP = 0.5 +# Inserted rows against the reference's NVFP4 / E4M3 recipes: the kernel's RMSNorm sums the squares in another order +# and uses approximate rsqrt / reciprocals, so an element on a rounding tie can land one code apart. Measured over +# 46 steps: no differing scale byte, at most 3 differing value codes in a step (1 in 300 of the inserted elements). +# Such a code moves one key's (or value's) element by an E2M1 step, up to 44 ulps of an output row, which is why +# the output is gated against the decode over the op's own caches rather than the reference's. +CODE_MISMATCH_MIN = 4 +CODE_MISMATCH_FRACTION = 0.005 + + +def check_insert(s: SimpleNamespace, before: tuple, ref) -> None: + """Every row the step does not insert is untouched; inserted rows match the reference's recipe.""" + kv_data0, kv_scale0, idx0 = before + pages = s.kv_data.shape[0] + mask = new_slot_mask(s, pages) + keep = ~mask + for plane in (0, 1): + assert torch.equal(s.kv_data[:, plane, 0][keep], kv_data0[:, plane, 0][keep]), ( + "a row past the insert moved" + ) + assert torch.equal(s.kv_scale[:, 0, 0][keep], kv_scale0[:, 0, 0][keep]) + assert torch.equal(s.idx_k.view(torch.uint8)[:, 0][keep], idx0.view(torch.uint8)[:, 0][keep]) + v_now = torch.stack([v_scale_rows(s.kv_scale[p, 1, 0]) for p in range(pages)]) + v_was = torch.stack([v_scale_rows(kv_scale0[p, 1, 0]) for p in range(pages)]) + v_ref = torch.stack([v_scale_rows(ref.kv_scale[p, 1, 0]) for p in range(pages)]) + assert torch.equal(v_now[keep], v_was[keep]), "a V scale outside the inserted rows moved" + assert torch.equal(s.kv_scale[:, 0, 0][mask], ref.kv_scale[:, 0, 0][mask]), "K block scales" + assert torch.equal(v_now[mask], v_ref[mask]), "V block scales (4x4 token-quad order)" + n = int(mask.sum()) * HEAD_DIM + for plane, scales in ((0, s.kv_scale[:, 0, 0][mask]), (1, v_now[mask])): + got = nvfp4_dequantize(s.kv_data[:, plane, 0][mask], scales) + want = nvfp4_dequantize(ref.kv_data[:, plane, 0][mask], scales) + assert int((got != want).sum()) <= max(CODE_MISMATCH_MIN, CODE_MISMATCH_FRACTION * n), ( + f"plane {plane} E2M1 codes" + ) + got = s.idx_k.view(torch.uint8)[:, 0][mask] + want = ref.idx_k.view(torch.uint8)[:, 0][mask] + assert int((got != want).sum()) <= max(CODE_MISMATCH_MIN, CODE_MISMATCH_FRACTION * n), ( + "index-K E4M3 codes" + ) + + +def decode_over_written_caches(s: SimpleNamespace, packed: torch.Tensor) -> torch.Tensor: + """The reference's decode (bf16 [T, 2048]) of ``packed``'s queries over the caches as the op left them.""" + q = query_rows(packed, s.position_ids, s.q_norm_weight, s.eps, s.rotary_cos_sin) + return attend_cached( + q, + s.kv_data, + s.kv_scale, + s.block_table, + s.seq_lens, + s.k_global_scale, + s.v_global_scale, + s.sm_scale, + s.decode_query_len, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/_m3_nvfp4_cache.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/attention/test_modeling_v2_m3_attn_c1.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_c1.py new file mode 100644 index 000000000000..61f59e77ced0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_c1.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_attn_c1 entry (fused decode attention + o_proj GEMV in one 96-CTA launch).""" + +import inspect +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels.m3_attn_mega_c1 import ( # noqa: E501 + MegaWorkspace, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_attn_c1 import ( # noqa: E501 + m3_attn_c1 as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_fused_decode_attn import ( # noqa: E501 + mx_gemm, + mxfp8_rows, + unswizzle_scales, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_fused_attn_step import ( # noqa: E402 + ATTN_ELEMENT_ULP, + ATTN_RMS_ULP, + attention_args, + build_step, + check_insert, + decode_over_written_caches, + mxfp8_operand, +) +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_attn_c1 requires a CUDA device" + +QLEN = 6 + +# Attention and the cache insert: the bands of m3_fused_decode_attn (the same kernel code). +# o_proj GEMV, on the op's own MXFP8 hand-off: exact operands, fp32 split-K accumulation over 2048 products against +# an fp64 sum, one bf16 rounding on each side: at most one bf16 step apart, which is 2 of these ulps (2^-8 of the row +# scale) in the row maximum's binade. Measured: at most 1.12 ulps. +_GEMV_ELEMENT_ULP = 2.0 +_GEMV_RMS_ULP = 0.5 + + +def _args(s, ws, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + wo, wosf = mxfp8_operand(6144, 2048, 120, 122, g, 6144) + out = torch.empty(s.num_tokens, 2048, dtype=torch.bfloat16, device="cuda") + y = torch.empty(s.num_tokens, 6144, dtype=torch.bfloat16, device="cuda") + return [ + s.packed, + ws.flags, + *attention_args(s), + 1, + 1, + out, + ws.q8, + ws.sf, + True, + wo, + wosf, + y, + 0, + 0, + 0, + 1, + 8, + 0, + ] + + +def _arg(name): + """Position of a parameter in the op's argument list (``is_valid`` mirrors ``__call__``).""" + return list(inspect.signature(type(op).is_valid).parameters)[1:].index(name) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 -- the fixture from _validating + certifying(op) + seed = zlib.crc32(cell.why.encode()) % 2**31 + s = build_step([cell.spec["seq_len"]], QLEN, seed=seed) + ws = MegaWorkspace(torch.device("cuda")) + for _ in range(cell.spec["calls"]): + args = _args(s, ws, seed + 1) + out, y = args[_arg("output")], args[_arg("y")] + wo, wosf = args[_arg("o_weight")], args[_arg("o_weight_scale")] + parity = int(ws.flags[32]) & 1 + before = (s.kv_data.clone(), s.kv_scale.clone(), s.idx_k.clone()) + ref = op.reference(*args) + op(*args) + check_insert(s, before, ref) + want = decode_over_written_caches(s, s.packed) + assert_within_ulp(out.view(-1, 128), want.view(-1, 128), ATTN_ELEMENT_ULP, ATTN_RMS_ULP) + # The o_proj units read the attention's MXFP8 hand-off, which is the merge's recipe on the bf16 output. + q8 = ws.q8.view(2, 8, 2048)[parity, :QLEN] + ex = unswizzle_scales(ws.sf.view(2, -1)[parity], QLEN, 64) + q8_want, ex_want = mxfp8_rows(out) + assert torch.equal(q8, q8_want) and torch.equal(ex, ex_want), "MXFP8 hand-off" + y_staged = mx_gemm(q8, ex, wo.view(torch.uint8), unswizzle_scales(wosf, 6144, 64)) + assert_within_ulp(y.view(-1, 128), y_staged.view(-1, 128), _GEMV_ELEMENT_ULP, _GEMV_RMS_ULP) + assert bool((ws.q8.view(2, -1)[parity ^ 1] == 255).all()), ( + "the next call's buffer is at rest" + ) + assert int(ws.flags[32]) & 1 == parity ^ 1, "the call counter advanced" + + +def test_guard_refuses_a_handoff_not_at_rest() -> None: + """The o_proj units poll the hand-off for codes other than 0xFF: stale codes read as this call's.""" + s = build_step([300], QLEN, seed=5) + ws = MegaWorkspace(torch.device("cuda")) + ws.q8.zero_() + args = _args(s, ws, 6) + with pytest.raises(AssertionError, match="not at rest"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_slot_the_block_table_does_not_give() -> None: + s = build_step([300], QLEN, seed=5) + args = _args(s, MegaWorkspace(torch.device("cuda")), 6) + loc = s.out_cache_loc.clone() + loc[2] = (loc[2] + 7) % (s.kv_data.shape[0] * 128) + args[_arg("out_cache_loc")] = loc + with pytest.raises(AssertionError, match="out_cache_loc must be the slot"): + with validating(op): + op(*args) + + +def test_op_rejects_another_query_length() -> None: + """Compiled for one request of 6 query tokens (96 CTAs, one o_proj unit each) only.""" + s = build_step([300], 4, seed=5) + args = _args(s, MegaWorkspace(torch.device("cuda")), 6) + args[_arg("decode_query_len")] = 4 + with pytest.raises(RuntimeError, match="six query tokens"): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_mega.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_mega.py new file mode 100644 index 000000000000..b5a48b0252eb --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_mega.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_attn_mega entry (QKV GEMV + fused decode attention + o_proj GEMV in one launch).""" + +import inspect +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels.m3_attn_mega import ( # noqa: E501 + MegaWorkspace, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_attn_mega import ( # noqa: E501 + m3_attn_mega as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_fused_decode_attn import ( # noqa: E501 + mx_gemm, + mxfp8_rows, + reference_attention, + unswizzle_scales, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_fused_attn_step import ( # noqa: E402 + ATTN_ELEMENT_ULP, + ATTN_RMS_ULP, + attention_args, + build_step, + check_insert, + decode_over_written_caches, + mxfp8_operand, +) +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_attn_mega requires a CUDA device" + +# Each stage is gated against its own reference fed the op's own input to that stage, so a band states one +# stage's error, not the propagation of an earlier stage's rounding through a quantizer. +# +# QKV GEMV and o_proj GEMV: exact MXFP8 operands, fp32 accumulation over 6144 / 2048 products in split-K partials +# against an fp64 sum, one bf16 rounding on each side: at most one bf16 step apart, which is 2 of these ulps (2^-8 +# of the row scale) in the row maximum's binade. Measured: bit-exact but for a few elements (0.43 ulp at most). +_GEMV_ELEMENT_ULP = 2.0 +_GEMV_RMS_ULP = 0.5 +# Attention and the cache insert: the bands of m3_fused_decode_attn (the same kernel code), on the op's own rows. + + +def _operands(seed): + g = torch.Generator(device="cuda").manual_seed(seed) + xq, xsf = mxfp8_operand(1, 6144, 126, 128, g, 128) + w, wsf = mxfp8_operand(2560, 6144, 119, 121, g, 2560) + wo, wosf = mxfp8_operand(6144, 2048, 120, 122, g, 6144) + return xq, xsf, w, wsf, wo, wosf + + +def _args(s, ops, ws): + xq, xsf, w, wsf, wo, wosf = ops + out = torch.empty(1, 2048, dtype=torch.bfloat16, device="cuda") + y = torch.empty(1, 6144, dtype=torch.bfloat16, device="cuda") + return [ + xq, + xsf, + w, + wsf, + ws.qkv_out, + ws.flags, + *attention_args(s), + 1, + 1, + out, + ws.q8, + ws.sf, + True, + wo, + wosf, + y, + 0, + 0, + 0, + 1, + 8, + 0, + ] + + +def _arg(name): + """Position of a parameter in the op's argument list (``is_valid`` mirrors ``__call__``).""" + return list(inspect.signature(type(op).is_valid).parameters)[1:].index(name) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 -- the fixture from _validating + certifying(op) + seed = zlib.crc32(cell.why.encode()) % 2**31 + s = build_step([cell.spec["seq_len"]], 1, seed=seed) + ops = _operands(seed + 1) + wo, wosf = ops[4], ops[5] + ws = MegaWorkspace(torch.device("cuda")) + for _ in range(cell.spec["calls"]): + args = _args(s, ops, ws) + out, y = args[_arg("output")], args[_arg("y")] + parity = int(ws.flags[32]) & 1 + caches = (s.kv_data.clone(), s.kv_scale.clone(), s.idx_k.clone()) + ref = op.reference(*args) + op(*args) + packed = ws.qkv_out[parity, :1] + assert_within_ulp( + packed.view(-1, 128), ref.packed.view(-1, 128), _GEMV_ELEMENT_ULP, _GEMV_RMS_ULP + ) + check_insert(s, caches, reference_attention(packed, *caches, *attention_args(s)[3:])) + want = decode_over_written_caches(s, packed) + assert_within_ulp(out.view(-1, 128), want.view(-1, 128), ATTN_ELEMENT_ULP, ATTN_RMS_ULP) + # The o_proj units read the attention's MXFP8 hand-off, which is the merge's recipe on the bf16 output. + q8 = ws.q8.view(2, 8, 2048)[parity, :1] + ex = unswizzle_scales(ws.sf.view(2, -1)[parity], 1, 64) + q8_want, ex_want = mxfp8_rows(out) + assert torch.equal(q8, q8_want) and torch.equal(ex, ex_want), "MXFP8 hand-off" + y_staged = mx_gemm(q8, ex, wo.view(torch.uint8), unswizzle_scales(wosf, 6144, 64)) + assert_within_ulp(y.view(-1, 128), y_staged.view(-1, 128), _GEMV_ELEMENT_ULP, _GEMV_RMS_ULP) + # The next call's buffers are back at rest. + assert bool((ws.qkv_out.view(torch.int16)[parity ^ 1] == -1).all()) + assert bool((ws.q8.view(2, -1)[parity ^ 1] == 255).all()) + assert int(ws.flags[32]) & 1 == parity ^ 1, "the call counter advanced" + + +def test_guard_refuses_a_workspace_not_at_rest() -> None: + """The attention CTAs poll the QKV rows for the 0xFFFF marker: stale rows read as this call's.""" + s = build_step([300], 1, seed=5) + ws = MegaWorkspace(torch.device("cuda")) + ws.qkv_out.zero_() + args = _args(s, _operands(6), ws) + with pytest.raises(AssertionError, match="not at rest"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_request_past_the_identity_window() -> None: + s = build_step([300], 1, seed=5) + args = _args(s, _operands(6), MegaWorkspace(torch.device("cuda"))) + seq_lens = s.seq_lens.clone() + seq_lens[0] = 2049 + args[_arg("seq_lens")] = seq_lens + with pytest.raises(AssertionError, match="identity selection covers 16 blocks"): + with validating(op): + op(*args) + + +def test_op_rejects_more_than_one_token() -> None: + """Compiled for one request with one query token only.""" + s = build_step([300, 400], 1, seed=5) + args = _args(s, _operands(6), MegaWorkspace(torch.device("cuda"))) + args[_arg("xq")] = torch.cat([args[_arg("xq")]] * 2) + args[_arg("output")] = torch.empty(2, 2048, dtype=torch.bfloat16, device="cuda") + args[_arg("y")] = torch.empty(2, 6144, dtype=torch.bfloat16, device="cuda") + with pytest.raises(RuntimeError, match="one request with one query token"): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_dense_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_dense_decode.py new file mode 100644 index 000000000000..5a7f07733314 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_dense_decode.py @@ -0,0 +1,149 @@ +# 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_103__tp4 import ( # noqa: E501 + dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__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) + + +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_103__tp4/attention/test_modeling_v2_m3_fused_decode_attn.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_fused_decode_attn.py new file mode 100644 index 000000000000..4db2720f9564 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_fused_decode_attn.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fused_decode_attn entry (fused norm/RoPE/KV-insert + identity sparse decode + merge).""" + +import inspect +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels.m3_fused_attn import ( # noqa: E501 + default_geometry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_fused_decode_attn import ( # noqa: E501 + PAGE, + mxfp8_rows, + unswizzle_scales, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_fused_decode_attn import ( # noqa: E501 + m3_fused_decode_attn as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_fused_attn_step import ( # noqa: E402 + ATTN_ELEMENT_ULP, + ATTN_RMS_ULP, + MXFP8_ROWS, + attention_args, + build_step, + check_insert, + decode_over_written_caches, +) +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_fused_decode_attn requires a CUDA device" + + +def _buffers(s): + t = s.num_tokens + out = torch.empty(t, 2048, dtype=torch.bfloat16, device="cuda") + q8 = torch.zeros(MXFP8_ROWS, 2048, dtype=torch.float8_e4m3fn, device="cuda") + sf = torch.zeros(MXFP8_ROWS * 64, dtype=torch.uint8, device="cuda") + return out, q8, sf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 -- the fixture from _validating + """The insert against the reference's recipe, the output against the reference's decode over the caches the + op wrote (see ATTN_ELEMENT_ULP / CODE_MISMATCH_FRACTION), the MXFP8 hand-off bit-exactly.""" + certifying(op) + spec = cell.spec + qg, bpc = spec["geometry"] + assert default_geometry(spec["qlen"], len(spec["seq_lens"])) == (qg, bpc), ( + "the cell's geometry is the target's" + ) + s = build_step(spec["seq_lens"], spec["qlen"], seed=zlib.crc32(cell.why.encode()) % 2**31) + out, q8, sf = _buffers(s) + args = (s.packed, *attention_args(s), qg, bpc, out, q8, sf, True) + before = (s.kv_data.clone(), s.kv_scale.clone(), s.idx_k.clone()) + ref = op.reference(*args) + op(*args) + check_insert(s, before, ref) + want = decode_over_written_caches(s, s.packed) + assert_within_ulp(out.view(-1, 128), want.view(-1, 128), ATTN_ELEMENT_ULP, ATTN_RMS_ULP) + # The MXFP8 hand-off is the merge's recipe applied to the op's own bf16 rows: bit-exact. + t = s.num_tokens + q8_want, ex_want = mxfp8_rows(out) + assert torch.equal(q8[:t].view(torch.uint8), q8_want), "MXFP8 hand-off values" + assert torch.equal(unswizzle_scales(sf, t, 64), ex_want), "MXFP8 hand-off scales" + + +def _cell_inputs(index=1): + spec = op.CELLS[index].spec + s = build_step(spec["seq_lens"], spec["qlen"], seed=7) + out, q8, sf = _buffers(s) + qg, bpc = spec["geometry"] + return s, [s.packed, *attention_args(s), qg, bpc, out, q8, sf, True] + + +def _arg(args, name): + """Position of a parameter in the op's argument list (``is_valid`` mirrors ``__call__``).""" + return list(inspect.signature(type(op).is_valid).parameters)[1:].index(name) + + +def test_guard_refuses_a_request_past_the_identity_window() -> None: + """16 blocks of 128: a 2049-token request would have its new token placed in block 15 and its tail unread.""" + s, args = _cell_inputs() + seq_lens = s.seq_lens.clone() + seq_lens[0] = 16 * PAGE + 1 + args[_arg(args, "seq_lens")] = seq_lens + with pytest.raises(AssertionError, match="identity selection covers 16 blocks"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_slot_the_block_table_does_not_give() -> None: + """The attention reads the inserted page from the cache it loaded: a slot elsewhere is never attended.""" + s, args = _cell_inputs() + loc = s.out_cache_loc.clone() + loc[0] = (loc[0] + 1) % (s.kv_data.shape[0] * PAGE) + args[_arg(args, "out_cache_loc")] = loc + with pytest.raises(AssertionError, match="out_cache_loc must be the slot"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_block_table_narrower_than_the_length() -> None: + s, args = _cell_inputs() + args[_arg(args, "block_table")] = s.block_table[:, :1].clone() + with pytest.raises(AssertionError, match="block_table must hold every attended block"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_strided_kv_pool() -> None: + """The kernel reads each [128, 64] head block as contiguous bytes; a transposed view is read as if it were.""" + s, args = _cell_inputs() + args[_arg(args, "kv_data")] = s.kv_data.transpose(3, 4).contiguous().transpose(3, 4) + with pytest.raises(AssertionError, match="kv_data must be"): + with validating(op): + op(*args) + + +def test_op_rejects_an_uncompiled_geometry() -> None: + """Only the target's geometries are compiled; anything else is refused by the op itself.""" + s, args = _cell_inputs(index=0) + args[_arg(args, "tokens_per_cluster")] = 2 + with pytest.raises(RuntimeError, match="is not compiled"): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_index_score.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_index_score.py new file mode 100644 index 000000000000..341fb87a4993 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_index_score.py @@ -0,0 +1,92 @@ +# 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_103__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_103__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 + +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" + + +@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_103__tp4/attention/test_modeling_v2_m3_producer.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_producer.py new file mode 100644 index 000000000000..803ee09dfee3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_producer.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 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_103__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_103__tp4.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, 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)) + + +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_103__tp4/attention/test_modeling_v2_m3_select_long.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_select_long.py new file mode 100644 index 000000000000..42f804302ef7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_select_long.py @@ -0,0 +1,159 @@ +# 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_103__tp4.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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) + + +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_103__tp4/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_sparse_decode.py new file mode 100644 index 000000000000..a57127e9e44d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_sparse_decode.py @@ -0,0 +1,121 @@ +# 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 + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_sparse_decode import ( # noqa: E501 + m3_sparse_decode as op, +) + +__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 + +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) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py new file mode 100644 index 000000000000..21fec5fc2607 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py @@ -0,0 +1,237 @@ +# 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_103__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_103__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 + +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, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_msa_kv_update.py new file mode 100644 index 000000000000..c5a2ec3f5af6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/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_103__tp4.kernels.msa_kv_update import ( # noqa: E501 + msa_kv_update as hook, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__tp4.ops.attention.msa_kv_update import ( # noqa: E501 + msa_kv_update as op, +) + +__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_103__tp4/comm/_m3_draft_ar_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_draft_ar_norm_op_matrix.py new file mode 100644 index 000000000000..c782ac2b782c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_draft_ar_norm_op_matrix.py @@ -0,0 +1,380 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fast draft loop's layer boundaries (``spec_worker.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. + +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 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 _target(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """The draft loop's boundary: ``spec_worker.draft_ar_norm`` with the drafter's norm module.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.spec_worker import ( # noqa: E501 + draft_ar_norm, + ) + from tensorrt_llm._torch.modules.rms_norm import RMSNorm + + norm = RMSNorm(hidden_size=HIDDEN, eps=EPS, dtype=torch.bfloat16).cuda() + with torch.no_grad(): + norm.weight.copy_(w) + return draft_ar_norm(ar, x, residual.clone(), norm) + + +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) + 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 + + 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_103__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..1b749e0c7482 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,663 @@ +# 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_103__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) + _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_103__tp4.kernels.m3_mc_allreduce as _ar # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__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 + + 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_103__tp4/comm/_rank_job.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_rank_job.py new file mode 100644 index 000000000000..eeeb202ae093 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_rank_job.py @@ -0,0 +1,81 @@ +# 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 + +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())] + assert len(devices) >= WORLD_SIZE, ( + f"this entry is certified at world size {WORLD_SIZE}; only " + f"{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, and a source checkout is not necessarily installed. + # tests/unittest/_torch/modeling_v2/models/minimax_m3//comm -> repo root. + repo_root = Path(__file__).resolve().parents[8] + assert (repo_root / "tensorrt_llm").is_dir(), ( + f"expected the repo root at {repo_root}, found no tensorrt_llm/ there; " + f"this file moved without its parents[] index following" + ) + env["PYTHONPATH"] = os.pathsep.join(p for p in (str(repo_root), env.get("PYTHONPATH", "")) if p) + + 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_103__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py new file mode 100644 index 000000000000..d2a966c7d5c1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py @@ -0,0 +1,30 @@ +# 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 ``spec_worker.draft_ar_norm`` 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 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_103__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_103__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..8b8585dc2677 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,38 @@ +# 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 + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.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 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") + + +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_103__tp4/gemm/test_modeling_v2_m3_mx_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mx_gemv.py new file mode 100644 index 000000000000..88e3d562d94a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mx_gemv.py @@ -0,0 +1,86 @@ +# 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_103__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 + +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) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..8ed3fc6df4a3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/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_103__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_103__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 + +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_103__tp4/gemm/test_modeling_v2_m3_target_head.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_target_head.py new file mode 100644 index 000000000000..dd55be9a924b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_target_head.py @@ -0,0 +1,248 @@ +# 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_103__tp4.kernels import ( # noqa: E501 + m3_target_head as target_head_kernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__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 + +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_103__tp4/gemm/test_modeling_v2_mxfp8_low_m_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_mxfp8_low_m_gemm.py new file mode 100644 index 000000000000..ed99da197a8a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_mxfp8_low_m_gemm.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the mxfp8_low_m_gemm entry of the MiniMax-M3 tp4 target.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm.mxfp8_low_m_gemm import ( # noqa: E501 + mxfp8_low_m_gemm as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "mxfp8_low_m_gemm requires a CUDA device" + + +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 _build(spec, seed): + """A BF16 activation and an MXFP8 weight: E4M3 values of a unit normal, UE8M0 block scales 2^-9 .. 2^-3.""" + g = torch.Generator(device="cuda").manual_seed(seed) + m, n, k = spec["m"], spec["n"], spec["k"] + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + w = torch.randn(n, k, generator=g, device="cuda").to(torch.float8_e4m3fn) + sf = torch.randint(118, 125, (n, k // 32), generator=g, device="cuda", dtype=torch.int32) + return x, w, _swizzle(sf.to(torch.uint8)) + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Accumulation order only, then one BF16 rounding. + + Both sides multiply the same E4M3 operands and power-of-two scales; the kernel sums each 32-wide block in the + MMA, rescales it exactly and accumulates blocks per warp, warps in order and split-K ranks in order, the + reference sums in a single fp32 GEMM. The fp32 results differ by a few fp32 ulps of the row's magnitude, which + moves the BF16 rounding by at most one step of the element -- at most 2 units of ``ULP[bf16]`` against the row + scale. Measured on GB300 over every cell: worst element 0.32 unit, relative RMS 0.006 unit. + """ + assert_within_ulp(out, ref, element_ulp=2.0, rms_ulp=0.25) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + x, w, wsf = _build(spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + args = (x, w, wsf, spec["warps"], spec["cluster_k"], spec["row_tiles"], False) + with validating(op): + out = op(*args) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + _compare(out, op.reference(*args)) + + +def test_the_gate_sees_a_wrong_scale_layout() -> None: + """Block scales read in the wrong order land far outside the band: most outputs pick up other blocks' + scales, which a layout bug in either the op or the reference would do.""" + spec = op.CELLS[1].spec + g = torch.Generator(device="cuda").manual_seed(7) + x, w, wsf = _build(spec, seed=7) + scrambled = wsf[torch.randperm(wsf.numel(), generator=g, device="cuda")] + args = (x, w, wsf, spec["warps"], spec["cluster_k"], spec["row_tiles"], False) + out = op(*args) + with pytest.raises(AssertionError): + _compare(out, op.reference(x, w, scrambled, *args[3:])) + + +def test_guard_refuses_a_short_scale_buffer() -> None: + """The kernel bulk-copies scale chunks by computed offset; a linear-length buffer is read past its end.""" + spec = op.CELLS[0].spec + x, w, wsf = _build(spec, seed=11) + n, k = w.shape + with pytest.raises(AssertionError, match="swizzled UE8M0 buffer"): + with validating(op): + op( + x, + w, + wsf[: n * k // 32 - 512], + spec["warps"], + spec["cluster_k"], + spec["row_tiles"], + False, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/_m3_moe_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/_m3_moe_operands.py new file mode 100644 index 000000000000..98a592728241 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/moe/test_modeling_v2_m3_moe16_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_moe16_engine.py new file mode 100644 index 000000000000..c9fe97ffd9f1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_moe16_engine.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 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_103__tp4.ops.moe.m3_moe16_engine import ( # noqa: E501 + M3Moe16State, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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 + +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_103__tp4/moe/test_modeling_v2_m3_moe_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_moe_engine.py new file mode 100644 index 000000000000..0b54f5fb2700 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_moe_engine.py @@ -0,0 +1,113 @@ +# 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_103__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_103__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 + +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) + + +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_103__tp4/moe/test_modeling_v2_m3_router_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_gemv.py new file mode 100644 index 000000000000..baa92e7c7522 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_gemv.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the router entries of the MiniMax-M3 tp4 target: m3_router_gemv, m3_router_gemv_fp4 and +m3_router_tc (one kernel source, one test module).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.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_103__tp4.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_103__tp4.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_103__tp4.ops.moe.m3_router_tc import ( # noqa: E501 + m3_router_tc as tc, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "the router entries require a CUDA device" + +#: 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" + + +@pytest.mark.parametrize("cell", gemv.CELLS, ids=[c.why[:48] for c in gemv.CELLS]) +def test_gemv_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], spec["e"], spec["k"], _seed(cell)) + with validating(gemv): + out = gemv(x, w, 1, False) + _compare_logits(out, gemv.reference(x, w, 1, False)) + + +@pytest.mark.parametrize("cell", gemv_fp4.CELLS, ids=[c.why[:48] for c in gemv_fp4.CELLS]) +def test_gemv_fp4_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], spec["e"], spec["k"], _seed(cell)) + sf_scale = torch.tensor([spec["sf_scale"]], dtype=torch.float32, device="cuda") + with validating(gemv_fp4): + logits, q, sf = gemv_fp4(x, w, sf_scale, False) + ref_logits, _, _ = gemv_fp4.reference(x, w, sf_scale, False) + _compare_logits(logits, ref_logits) + _compare_fp4(q, sf, x, sf_scale) + + +@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_the_gemv_is_a_true_fp32_product() -> None: + """And the GEMV, by contrast, is no TF32 product: the TF32 reference rejects it.""" + x, w = _build(1, 128, 6144, 6) + out = gemv(x, w, 1, False) + with pytest.raises(AssertionError): + _compare_logits(out, tc.reference(x, w, None, False)[0]) + + +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) + + +def test_guard_refuses_a_zero_global_scale() -> None: + x, w = _build(1, 128, 6144, 4) + with pytest.raises(AssertionError, match="positive and finite"): + with validating(gemv_fp4): + gemv_fp4(x, w, torch.zeros(1, dtype=torch.float32, device="cuda"), False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py new file mode 100644 index 000000000000..8778dbe2805d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py @@ -0,0 +1,71 @@ +# 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_103__tp4.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + VOCAB, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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 + +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_103__tp4/spec/_m3_spec_inputs.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/_m3_spec_inputs.py new file mode 100644 index 000000000000..5c5d4161aa78 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/spec/test_modeling_v2_m3_capture_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_capture_norm.py new file mode 100644 index 000000000000..52fdac496253 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/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_103__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 + +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_103__tp4/spec/test_modeling_v2_m3_chunk_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_chunk_norm.py new file mode 100644 index 000000000000..985b51e8650c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/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_103__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 + +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_103__tp4/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_input.py new file mode 100644 index 000000000000..ca13b3efb063 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/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_103__tp4.kernels import ( # noqa: E501 + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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 + +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_103__tp4/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_step.py new file mode 100644 index 000000000000..e6c2b67bf6c0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_step.py @@ -0,0 +1,686 @@ +# 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_103__tp4.kernels import ( # noqa: E501 + m3_draft_step as draft_step_kernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__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_103__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_103__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 + +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 (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): + 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: + """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) + 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: + """DRAFT_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.dispatch, "DRAFT_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) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py new file mode 100644 index 000000000000..f66014d47aac --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py @@ -0,0 +1,293 @@ +# 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_103__tp4 import ( # noqa: E501 + linears, + spec_worker, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__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 ( + spec_worker.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(spec_worker.M3Eagle3Worker, name) is getattr(Eagle3OneModelWorker, name), ( + name + ) + for module, name in ( + (spec_worker, "M3DraftGemvLinearMethod"), + (spec_worker, "adopt_draft_linears"), + (spec_worker, "greedy_head_outputs"), + (linears, "M3FP8LMHeadMethod"), + (linears, "adopt_fp8_lm_heads"), + ): + assert not hasattr(module, name), f"{module.__name__}.{name}" + pkg = spec_worker.__name__.rsplit(".", 1)[0] + for gone in ("kernels.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_103__tp4/spec/test_modeling_v2_m3_sampled_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_sampled_accept.py new file mode 100644 index 000000000000..189d28ac638f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_sampled_accept.py @@ -0,0 +1,167 @@ +# 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_103__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_103__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 + +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) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_shard_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_shard_accept.py new file mode 100644 index 000000000000..a984ca817a2e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/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_103__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_103__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 + +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_103__tp4/spec/test_modeling_v2_m3_slot_scatter.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_slot_scatter.py new file mode 100644 index 000000000000..ec99ade9b043 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/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_103__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 + +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_103__tp4/spec/test_modeling_v2_m3_step_gather.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_step_gather.py new file mode 100644 index 000000000000..d75053d8c674 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/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_103__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 + +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_103__tp4/test_modeling_v2_m3_capture_boundary.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_capture_boundary.py new file mode 100644 index 000000000000..8d8835ad067e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_capture_boundary.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decoder layer's feed-forward exit on Eagle3 capture layers (``M3DecoderLayer._feed_forward_boundary``). + +A capture layer with a next norm keeps the fused boundary (all-reduce + residual + next norm in one kernel) and +captures the reduced residual sum it writes; the last layer captures its raw output with the residual, as the +in-tree capture does. CPU only: the boundary and the capture are stand-ins. +""" + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4 import ( # noqa: E501 + layers, +) + + +class _Layer: + """The attributes the boundary reads, with a recording capture and next-norm stand-in.""" + + _fused_feed_forward_boundary = layers.M3DecoderLayer._fused_feed_forward_boundary + _feed_forward_boundary = layers.M3DecoderLayer._feed_forward_boundary + + def __init__(self, has_next_norm: bool): + self.layer_idx = 56 + self.next_layer_layernorm = object() if has_next_norm else None + self.captured = [] + self.boundaries = [] + + def capture(self, spec_metadata, layer_id, hidden_states, residual): + self.captured.append((layer_id, hidden_states, residual)) + + def _next_layer_norm(self, hidden_states, residual, fused_boundary): + self.boundaries.append(fused_boundary) + if self.next_layer_layernorm is None: + return hidden_states, residual + return ("normed", hidden_states), ("updated", hidden_states, residual) + + +def test_a_capture_layer_keeps_the_fused_boundary_and_captures_its_residual_sum() -> None: + layer = _Layer(has_next_norm=True) + fused = layer._fused_feed_forward_boundary(capture=True) + assert fused + h, r = torch.ones(2), torch.zeros(2) + out = layer._feed_forward_boundary(h, r, fused, True, None) + assert layer.boundaries == [True] + assert layer.captured == [(56, ("updated", h, r), None)] + assert out == (("normed", h), ("updated", h, r)) + + +def test_the_last_layer_captures_its_raw_output_with_the_residual() -> None: + layer = _Layer(has_next_norm=False) + fused = layer._fused_feed_forward_boundary(capture=True) + assert not fused + h, r = torch.ones(2), torch.zeros(2) + out = layer._feed_forward_boundary(h, r, fused, True, None) + assert layer.captured == [(56, h, r)] and layer.boundaries == [False] + assert out == (h, r) + + +def test_a_layer_without_capture_is_unchanged() -> None: + layer = _Layer(has_next_norm=True) + assert layer._fused_feed_forward_boundary(capture=False) + h, r = torch.ones(2), torch.zeros(2) + layer._feed_forward_boundary(h, r, True, False, None) + assert not layer.captured and layer.boundaries == [True] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_dispatch.py new file mode 100644 index 000000000000..7a06b668d5a4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_dispatch.py @@ -0,0 +1,836 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode target's dispatch table (``dispatch.py``): its row ranges against the catalog cells that certify +them and the kernels' own limits, the selection for every decode row count, and the miss / uncovered policy. + +CPU only: nothing here launches a kernel. +""" + +import ast +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4 import ( # noqa: E501 + dispatch, + layers, + spec_worker, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 + m3_attn_mega_c1, + m3_draft_step, + m3_fused_attn, + m3_mc_allreduce, + m3_shard_accept, + m3_target_head, + triton_sparse_decode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 + mxfp8_low_m_gemm as low_m, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.modeling import ( # noqa: E501 + DecodeTarget, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm import ( # noqa: E501 + m3_target_head as target_head_entry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm import ( # noqa: E501 + mxfp8_low_m_gemm as low_m_entry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm.m3_mx_gemv import ( # noqa: E501 + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + SHARED_DOWN, + SHARED_GATE_UP, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec import ( # noqa: E501 + m3_draft_step as draft_step_entry, +) +from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.mx_gemv import op as mx_gemv +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.linear import TensorParallelMode + +PROJECTIONS = (QKV_SPARSE, QKV, O_PROJ, SHARED_GATE_UP, SHARED_DOWN, DENSE_GATE_UP, DENSE_DOWN) +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() + + +class _Log: + def __init__(self): + self.info, self.warning = [], [] + + +@pytest.fixture +def log(monkeypatch) -> _Log: + """The dispatch's log lines, by level.""" + rec = _Log() + monkeypatch.setattr(dispatch.logger, "info", rec.info.append) + monkeypatch.setattr(dispatch.logger, "warning", rec.warning.append) + return rec + + +# ------------------------------------------------------------------ the ranges against their certificates + + +@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)) + + +# ------------------------------------------------------------------ the ranges against the kernels' own limits + + +@pytest.mark.parametrize("nk", PROJECTIONS, ids=str) +def test_gemv_rows_are_the_rows_the_kernel_serves(nk) -> None: + n, k = nk + served = {m for m in range(1, mx_gemv.MAX_TOKENS + 1) if mx_gemv.serves(n, k, m, False)} + assert dispatch.MX_GEMV_ROWS[nk] == served + served_bf16 = [m for m in range(1, mx_gemv.MAX_TOKENS + 1) if mx_gemv.serves(n, k, m, True)] + assert served_bf16 == list(range(1, dispatch.MX_GEMV_BF16_MAX_TOKENS[nk] + 1)) + + +@pytest.mark.parametrize("nk", sorted(dispatch.LOW_M_GEMM_MAX_TOKENS), ids=str) +def test_low_m_gemm_rows_are_its_configuration_limit(nk) -> None: + warps, cluster_k, max_m, row_tiles = low_m.get_config(*nk) + assert dispatch.LOW_M_GEMM_MAX_TOKENS[nk] == max_m + for m in range(1, max_m + 1): + assert low_m._smem_bytes(m, nk[1], warps, cluster_k, row_tiles) <= low_m._MAX_DYNAMIC_SMEM + + +def test_the_low_m_gemm_has_a_range_for_every_configuration() -> None: + assert set(dispatch.LOW_M_GEMM_MAX_TOKENS) == set(low_m._CONFIGS) + + +def test_c1_takes_exactly_its_range() -> None: + for qlen in range(1, dispatch.FUSED_ATTN_MAX_QUERY_LEN + 1): + for reqs in range(1, dispatch.FUSED_ATTN_MAX_TOKENS + 1): + assert m3_attn_mega_c1.c1_applies(qlen, reqs) == dispatch.ATTN_C1.covers((qlen, reqs)) + + +def test_every_fused_attention_shape_has_a_compiled_geometry() -> None: + for qlen, reqs in dispatch.FUSED_DECODE_ATTN.keys: + assert (qlen, *m3_fused_attn.default_geometry(qlen, reqs)) in m3_fused_attn.GEOMETRIES + + +def test_the_plans_are_the_all_reduce_buckets() -> None: + assert tuple(p.max_tokens for p in dispatch.DECODE_PLANS) == m3_mc_allreduce.BUCKETS + assert dispatch.DECODE_MAX_TOKENS == m3_mc_allreduce.MAX_TOKENS == mx_gemv.MAX_TOKENS + + +def test_the_thresholds_are_the_kernels() -> None: + assert dispatch.SPARSE_DECODE_MQ_MIN_TOKENS == triton_sparse_decode.MQ_MIN_TOKENS + assert dispatch.MOE16_MAX_TOKENS < dispatch.ATTN_PUSH_MAX_TOKENS <= dispatch.DECODE_MAX_TOKENS + verify_rows = dispatch.SHARD_ACCEPT_MAX_REQUESTS * (dispatch.SHARD_ACCEPT_DRAFT_LEN + 1) + assert verify_rows == dispatch.DECODE_MAX_TOKENS == m3_shard_accept.SHARD_ROWS + + +def test_no_fast_path_quantizes_a_bf16_head_or_the_drafter() -> None: + """The target carries the exact-numerics configuration only: the lm_head and the Eagle3 drafter run at the + checkpoint's precision, so no fast path (nor a range constant) of an FP8 / NVFP4 / pruned head or of a quantized + drafter projection is left to select.""" + names = {r.path for r in dispatch.FAST_PATHS} + assert not names & { + "fp8_lm_head", + "draft_head", + "m3_draft_gemv", + "m3_draft_in_gemv", + "m3_gemv_ar", + } + for constant in ( + "FP8_LM_HEAD", + "FP8_LM_HEAD_MAX_TOKENS", + "DRAFT_HEAD", + "DRAFT_HEAD_MAX_TOKENS", + ): + assert not hasattr(dispatch, constant), constant + + +# ------------------------------------------------------------------ selection, every decode row count + + +def _expected_moe(rows: int): + if rows <= 8: + return "m3_moe_engine" + if rows <= 13: + return "m3_moe16_engine" + if rows <= 48: + return "mc_finalize_allreduce_norm" + return None + + +def _expected_router(rows: int): + if rows == 1: + return "m3_router_gemv" + if rows <= 48: + return "m3_router_tc" + return None + + +@pytest.mark.parametrize("rows", ROWS) +def test_selection(rows) -> None: + assert dispatch.is_decode_step(rows) == (rows <= 48) + if rows <= 48: + assert dispatch.plan_for(rows).name == ( + "tiny" if rows <= 8 else "small" if rows <= 16 else "verify" + ) + else: + with pytest.raises(AssertionError, match="no decode plan"): + dispatch.plan_for(rows) + assert dispatch.moe_path(rows) == _expected_moe(rows) + assert dispatch.router_path(rows) == _expected_router(rows) + assert dispatch.EMBED_FIRSTNORM.covers(rows) == (rows <= 8) + assert dispatch.ATTN_PUSH.covers(rows) == (rows <= 16) + assert dispatch.MOE_PUSH.covers(rows) == (rows <= 8) + assert dispatch.MC_ALLREDUCE.covers(rows) == (rows <= 48) + + +@pytest.mark.parametrize("rows", ROWS) +def test_linear_selection(rows) -> None: + bf16_max = { + QKV_SPARSE: 16, + QKV: 16, + O_PROJ: 32, + SHARED_GATE_UP: 16, + SHARED_DOWN: 16, + DENSE_GATE_UP: 8, + DENSE_DOWN: 16, + } + for nk in PROJECTIONS: + prequant_gemv = rows <= 48 and not (nk == SHARED_DOWN and 17 <= rows <= 32) + assert dispatch.linear_path(*nk, rows, True, True) == ( + "m3_mx_gemv" if prequant_gemv else "flashinfer_mxfp8_gemm" + ) + for low_m_gemm in (False, True): + assert dispatch.linear_path(*nk, rows, False, low_m_gemm) == ( + "m3_mx_gemv_bf16" if rows <= bf16_max[nk] else "mxfp8_gemm" + ) + # A projection outside the table takes the generic GEMM at every row count. + assert dispatch.linear_path(4096, 4096, rows, True, True) == "flashinfer_mxfp8_gemm" + assert dispatch.linear_path(4096, 4096, rows, False, True) == "mxfp8_gemm" + + +def test_draft_step_selection() -> None: + """The fused draft loop takes the configured draft length at 1..8 requests (one head row per request, the + kernel's limit); its draft input every pass's rows: the first pass's 4 per request, then one per request.""" + assert ( + dispatch.DRAFT_STEP_MAX_REQUESTS + == draft_step_entry.MAX_TOKENS + == m3_draft_step.MAX_TOKENS + == 8 + ) + assert dispatch.DRAFT_STEP_DRAFT_LEN == dispatch.SHARD_ACCEPT_DRAFT_LEN + for draft_len in range(1, 6): + for nreq in range(1, 14): + covered = draft_len == 3 and nreq <= 8 + assert dispatch.DRAFT_STEP.covers((draft_len, nreq)) == covered + assert dispatch.DRAFT_INPUT.keys == frozenset(range(1, 9)) | {12, 16, 20, 24, 28, 32} + + +def test_verify_selection() -> None: + for draft_len in range(1, 6): + for nreq in range(1, 14): + covered = draft_len == 3 and nreq * (draft_len + 1) <= 48 + assert dispatch.SHARD_ACCEPT.covers((draft_len, nreq)) == covered + + +def test_target_head_selection() -> None: + """The coded target head takes every verify of at most 8 rows (one or two requests at draft length 3); the early + fc fork the same verifies.""" + assert dispatch.TARGET_HEAD_MAX_ROWS == 8 == dispatch.FC_EARLY_MAX_ROWS + for rows in ROWS: + assert dispatch.TARGET_HEAD.covers(rows) == (rows <= 8) + + +@pytest.mark.parametrize("rows", ROWS) +def test_attention_shapes(rows) -> None: + for qlen in range(1, 9): + if rows % qlen: + continue + reqs = rows // qlen + assert dispatch.FUSED_DECODE_ATTN.covers((qlen, reqs)) == (rows <= 16) + assert dispatch.ATTN_MEGA.covers((qlen, reqs)) == (qlen == 1 and reqs == 1) + assert dispatch.ATTN_C1.covers((qlen, reqs)) == (qlen == 6 and reqs == 1) + + +# ------------------------------------------------------------------ require mode and fallback logging + + +def test_a_miss_raises_under_require() -> None: + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_moe_engine'.* 5 rows.*not built"): + dispatch.missed("m3_moe_engine", 5, "the engine is not built") + + +def test_a_miss_logs_once_per_reason_otherwise(log) -> None: + for _ in range(3): + dispatch.missed("m3_moe_engine", 5, "the engine is not built") + dispatch.missed("m3_moe_engine", 6, "the engine failed to build") + assert len(log.warning) == 2 + assert "m3_moe_engine at 5 rows: the engine is not built" in log.warning[0] + + +def test_uncovered_never_raises_and_logs_once(log) -> None: + dispatch.set_require(True) + for _ in range(3): + dispatch.uncovered("fused_decode_attn", "a request beyond 16 KV blocks") + assert len(log.info) == 1 and not log.warning + + +class _Core: + """The parts of the text model ``DecodeTarget._embed`` reads: no embedding owners built.""" + + def __init__(self): + self._embed_owners = {} + + @staticmethod + def embed_tokens(input_ids: torch.Tensor) -> torch.Tensor: + return input_ids.to(torch.float32)[:, None].expand(-1, 4) + + +@pytest.mark.parametrize("rows", (1, 8)) +def test_embedding_without_its_owner_raises_under_require(rows) -> None: + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=rf"'embed_firstnorm'.* {rows} rows.*owner" + ): + DecodeTarget(_Core())._embed(torch.arange(rows, dtype=torch.int32), rows) + + +def test_embedding_without_its_owner_falls_back_and_logs_once(log) -> None: + target = DecodeTarget(_Core()) + for _ in range(2): + hidden, residual = target._embed(torch.arange(3, dtype=torch.int32), 3) + assert residual is None and hidden.shape == (3, 4) + assert len(log.warning) == 1 and "embed_firstnorm at 3 rows" in log.warning[0] + + +def test_embedding_past_its_rows_is_uncovered_in_every_mode(log) -> None: + dispatch.set_require(True) + hidden, residual = DecodeTarget(_Core())._embed(torch.arange(9, dtype=torch.int32), 9) + assert residual is None and hidden.shape == (9, 4) + assert len(log.info) == 1 and "embed_firstnorm" in log.info[0] and not log.warning + + +def test_a_push_without_its_workspace_raises_under_require(monkeypatch) -> None: + monkeypatch.setattr(m3_mc_allreduce, "get_workspace", lambda mapping, hidden: None) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'attn_push'.* 4 rows"): + DecodeTarget(_Core())._attn_push_target(_Layer(), 4, 6144) + + +def test_a_push_without_its_workspace_falls_back_and_logs_once(monkeypatch, log) -> None: + monkeypatch.setattr(m3_mc_allreduce, "get_workspace", lambda mapping, hidden: None) + target = DecodeTarget(_Core()) + assert target._attn_push_target(_Layer(), 4, 6144) is None + assert target._moe_push_targets(_Layer(), 4, 6144) is None + assert target._attn_push_target(_Layer(), 4, 6144) is None + assert len(log.warning) == 2 + assert "attn_push at 4 rows" in log.warning[0] and "moe_push at 4 rows" in log.warning[1] + + +class _Layer: + """The parts of a decoder layer the push targets read.""" + + mapping = None + + +def test_every_dispatch_name_the_target_uses_exists() -> None: + """A fallback message or selection that names a missing ``dispatch.X`` fails only when that branch runs.""" + target_dir = Path(dispatch.__file__).parent + missing = [] + for path in sorted(target_dir.rglob("*.py")): + for node in ast.walk(ast.parse(path.read_text())): + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "dispatch" + and not hasattr(dispatch, node.attr) + ): + missing.append( + f"{path.relative_to(target_dir)}:{node.lineno}: dispatch.{node.attr}" + ) + assert not missing, missing + + +def test_low_m_gemm_cells_follow_its_configurations() -> None: + for (n, k), (warps, cluster_k, row_tiles, max_m) in low_m_entry.CONFIGS.items(): + assert low_m.get_config(n, k) == (warps, cluster_k, max_m, row_tiles) + assert dispatch.LOW_M_GEMM_MAX_TOKENS[(n, k)] == max_m + + +# ------------------------------------------------------------------ the fused draft loop's preconditions + +_HEAD_ROWS = 64 +_HIDDEN = 32 + + +def _drafter(head_weight=None): + """A drafter with only what M3Eagle3Worker._m3_drafter_supported reads: one plain Eagle3 layer with a gated MLP, + a whole-table embedding, all-reduced o_proj and down_proj and a vocab-sharded BF16 head.""" + mlp = GatedMLP.__new__(GatedMLP) + torch.nn.Module.__init__(mlp) + mlp.down_proj = SimpleNamespace(all_reduce=object()) + layer = SimpleNamespace( + _next_layer_regular=False, + mlp=mlp, + self_attn=SimpleNamespace(o_proj=SimpleNamespace(all_reduce=object())), + ) + model = SimpleNamespace( + num_layers=1, + _use_mla=False, + _eh_proj_before_attn=False, + midlayer=layer, + embed_tokens=SimpleNamespace(tp_size=1), + ) + weight = ( + head_weight + if head_weight is not None + else torch.zeros((_HEAD_ROWS, _HIDDEN), dtype=torch.bfloat16) + ) + lm_head = SimpleNamespace( + weight=weight, bias=None, tp_mode=TensorParallelMode.COLUMN, padding_size=0, tp_size=4 + ) + return SimpleNamespace(model=model, lm_head=lm_head) + + +def _worker(): + """An M3Eagle3Worker past the step's verify: the fused verify left the first-pass gather rows and the fused + first draft inputs; nothing else of the in-tree worker is read.""" + worker = spec_worker.M3Eagle3Worker.__new__(spec_worker.M3Eagle3Worker) + worker.mapping = SimpleNamespace(tp_size=4, tp_rank=1, enable_attention_dp=False) + worker._d2t = None + worker._m3_draft_head = None + worker._m3_first_gather_ids = torch.zeros(1, dtype=torch.int64) + worker._m3_fused_first_inputs = True + return worker + + +_ATTN_MD = SimpleNamespace(kv_lens_cuda=None) +_SPEC_MD = SimpleNamespace(runtime_draft_len=3) + + +@pytest.fixture +def no_gpu_draft_head(monkeypatch): + """The fused head's CUDA-only lookups replaced: a 148-SM device, and a built exchange workspace.""" + monkeypatch.setattr(m3_draft_step, "_num_sms", lambda: 148) + monkeypatch.setattr(m3_draft_step, "_WS", object()) + + +def test_draft_loop_without_the_prebuilt_head_raises_under_require(no_gpu_draft_head) -> None: + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=r"'m3_draft_step'.* 4 rows.*prebuild_draft_head" + ): + _worker()._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, _drafter()) + + +def test_draft_loop_without_the_prebuilt_head_falls_back_and_logs_once( + no_gpu_draft_head, log +) -> None: + worker = _worker() + for _ in range(2): + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, _drafter()) + assert len(log.warning) == 1 and "m3_draft_step at 4 rows" in log.warning[0] + + +def test_draft_loop_with_a_head_built_for_another_weight_raises_under_require( + no_gpu_draft_head, +) -> None: + worker = _worker() + assert worker.prebuild_draft_head(_drafter()) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_draft_step'"): + worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 2, _drafter()) + + +def test_draft_loop_takes_a_covered_step_once_the_head_is_built(no_gpu_draft_head) -> None: + dispatch.set_require(True) + drafter = _drafter() + worker = _worker() + assert worker.prebuild_draft_head(drafter) + head = worker._m3_draft_head + assert head.weight is drafter.lm_head.weight and head.vocab_offset == _HEAD_ROWS + assert torch.equal(draft_step_entry.from_fragments(head.wf), drafter.lm_head.weight) + for nreq in range(1, 9): + assert worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, nreq, drafter) + + +def test_draft_loop_streams_the_coded_head_once_built(no_gpu_draft_head) -> None: + """A drafter head whose K the coding covers is prebuilt in its 12-bit coding (decoding to the weight bit for bit) + and serves every covered step; one whose coding fails its round trip keeps the BF16 fragments (fail closed).""" + dispatch.set_require(True) + g = torch.Generator().manual_seed(11) + weight = (torch.randn((_HEAD_ROWS, 128), generator=g) * 0.02).to(torch.bfloat16) + drafter = _drafter(weight) + worker = _worker() + assert worker.prebuild_draft_head(drafter) + head = worker._m3_draft_head + assert head.coded and head.weight is weight and head.vocab_offset == _HEAD_ROWS + decoded = draft_step_entry.head_weight( + head.wf, 128, head.exc_off, head.exc_pos, head.exc_w, head.exp_base + ) + assert torch.equal(decoded.view(torch.int16), weight.view(torch.int16)) + for nreq in range(1, 9): + assert worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, nreq, drafter) + + +def test_draft_head_that_fails_its_round_trip_keeps_the_bf16_fragments( + no_gpu_draft_head, monkeypatch +) -> None: + real = target_head_entry.encode_chunked + + def corrupt(weight, *args, **kwargs): + packed, base, exc_off, exc_pos, exc_w = real(weight) + packed = packed.clone() + packed[5] ^= 2 + return packed, base, exc_off, exc_pos, exc_w + + monkeypatch.setattr(m3_draft_step, "encode_chunked", corrupt) + dispatch.set_require(True) + weight = torch.ones((_HEAD_ROWS, 128), dtype=torch.bfloat16) + drafter = _drafter(weight) + worker = _worker() + assert worker.prebuild_draft_head(drafter) + head = worker._m3_draft_head + assert not head.coded and head.exc_off is None + assert torch.equal(draft_step_entry.from_fragments(head.wf), weight) + assert worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + + +def test_draft_loop_past_its_requests_is_uncovered_in_every_mode(no_gpu_draft_head, log) -> None: + dispatch.set_require(True) + drafter = _drafter() + worker = _worker() + assert worker.prebuild_draft_head(drafter) + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 9, drafter) + assert not worker._m3_fused_loop_ok(_ATTN_MD, SimpleNamespace(runtime_draft_len=2), 4, drafter) + assert len(log.info) == 2 and not log.warning + + +def test_draft_loop_needs_the_fused_first_inputs(no_gpu_draft_head, log) -> None: + """A step whose first draft inputs came from the in-tree path (or that the fused verify did not prepare) drafts + with the in-tree loop, silently: it is not a draft-step shape.""" + dispatch.set_require(True) + drafter = _drafter() + worker = _worker() + assert worker.prebuild_draft_head(drafter) + worker._m3_fused_first_inputs = False + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + worker._m3_fused_first_inputs, worker._m3_first_gather_ids = True, None + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + assert not log.info and not log.warning + + +def test_an_unsupported_drafter_is_uncovered(no_gpu_draft_head, log) -> None: + dispatch.set_require(True) + worker = _worker() + padded = _drafter() + padded.lm_head.padding_size = 16 + two_layers = _drafter() + two_layers.model.num_layers = 2 + for drafter in (padded, two_layers): + assert not worker.prebuild_draft_head(drafter) + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + assert log.info and not log.warning + + +def test_a_capture_without_the_exchange_workspace_raises_under_require(monkeypatch) -> None: + monkeypatch.setattr(m3_draft_step, "_num_sms", lambda: 148) + monkeypatch.setattr(m3_draft_step, "_WS", None) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + drafter = _drafter() + worker = _worker() + assert worker.prebuild_draft_head(drafter) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_draft_step'.*workspace.*captured"): + worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + + +# ------------------------------------------------------------------ the coded target head's preconditions + +_TARGET_ROWS = 64 +_TARGET_K = 128 + + +def _target_lm_head(weight=None): + if weight is None: + g = torch.Generator().manual_seed(7) + weight = (torch.randn(_TARGET_ROWS, _TARGET_K, generator=g) * 0.02).to(torch.bfloat16) + return SimpleNamespace( + weight=weight, bias=None, tp_mode=TensorParallelMode.COLUMN, padding_size=0, tp_size=4 + ) + + +@pytest.fixture +def no_gpu_target_head(monkeypatch): + monkeypatch.setattr(m3_target_head, "_num_sms", lambda: 148) + + +def _target_worker(): + worker = spec_worker.M3Eagle3Worker.__new__(spec_worker.M3Eagle3Worker) + worker._m3_target_head = None + return worker + + +def test_target_head_without_the_prebuilt_copy_raises_under_require() -> None: + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=r"'m3_target_head'.* 4 rows.*prebuild_target_head" + ): + _target_worker()._m3_coded_head(_target_lm_head(), 4) + + +def test_target_head_without_the_prebuilt_copy_falls_back_and_logs_once(log) -> None: + worker = _target_worker() + for _ in range(2): + assert worker._m3_coded_head(_target_lm_head(), 8) is None + assert len(log.warning) == 1 and "m3_target_head at 8 rows" in log.warning[0] + + +def test_target_head_takes_a_covered_verify_once_built(no_gpu_target_head) -> None: + """The prebuilt copy decodes to the head's weight bit for bit and serves every covered row count; a copy built + for another weight is a miss.""" + dispatch.set_require(True) + lm_head = _target_lm_head() + worker = _target_worker() + assert worker.prebuild_target_head(lm_head) + head = worker._m3_target_head + assert head.coded and head.weight is lm_head.weight and head.grid == 148 + decoded = target_head_entry.decode( + head.w, head.base, head.exc_off, head.exc_pos, head.exc_w, _TARGET_ROWS, _TARGET_K + ) + assert torch.equal(decoded.view(torch.int16), lm_head.weight.view(torch.int16)) + for rows in range(1, 9): + assert worker._m3_coded_head(lm_head, rows) is head + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_target_head'"): + worker._m3_coded_head(_target_lm_head(lm_head.weight.clone()), 4) + + +def test_target_head_past_its_rows_is_uncovered_in_every_mode(no_gpu_target_head, log) -> None: + dispatch.set_require(True) + lm_head = _target_lm_head() + worker = _target_worker() + assert worker.prebuild_target_head(lm_head) + for rows in (9, 12, 32, 48): + assert worker._m3_coded_head(lm_head, rows) is None + assert log.info and not log.warning + + +def test_target_head_that_fails_its_round_trip_keeps_the_bf16_gemm( + no_gpu_target_head, monkeypatch, log +) -> None: + """Fail closed: a coding that does not decode to the weight bit for bit is dropped at prebuild (a warning), and the + verify runs the in-tree GEMM without a miss, in every mode.""" + real = target_head_entry.encode_chunked + + def corrupt(weight, *args, **kwargs): + packed, base, exc_off, exc_pos, exc_w = real(weight) + packed = packed.clone() + packed[3] ^= 1 + return packed, base, exc_off, exc_pos, exc_w + + monkeypatch.setattr(m3_target_head, "encode_chunked", corrupt) + dispatch.set_require(True) + lm_head = _target_lm_head() + worker = _target_worker() + assert not worker.prebuild_target_head(lm_head) + assert not worker._m3_target_head.coded + assert len(log.warning) == 1 and "bit for bit" in log.warning[0] + assert worker._m3_coded_head(lm_head, 4) is None + assert any("round trip" in line for line in log.info) + + +def test_an_unsupported_target_head_is_not_coded(no_gpu_target_head) -> None: + worker = _target_worker() + padded = _target_lm_head() + padded.padding_size = 16 + narrow = _target_lm_head(torch.zeros((_TARGET_ROWS, 96), dtype=torch.bfloat16)) + for lm_head in (padded, narrow): + assert not worker.prebuild_target_head(lm_head) + assert worker._m3_target_head is None + + +# ------------------------------------------------------------------ the early fc fork + + +class _CaptureMetadata: + def __init__(self, nreq: int): + self.layers_to_capture = (1, 29, 56) + self.num_capture_tokens = nreq * 4 + self.runtime_draft_len = 3 + self.captured = [] + + def maybe_capture_hidden_states(self, layer_id, hidden_states, residual) -> None: + self.captured.append(layer_id) + + +def _fc_early_worker(monkeypatch, eligible: bool = True): + worker = spec_worker.M3Eagle3Worker.__new__(spec_worker.M3Eagle3Worker) + worker._m3_capture_norm = None + worker._m3_fc_norm_w = None + worker._m3_fc_early_drafter = None + worker._m3_fc_early_stream = None + forks = [] + monkeypatch.setattr(worker, "tail_eligible", lambda *args: eligible) + + def fork(md, num_tokens, draft_model, after=None): + forks.append((num_tokens, draft_model, after)) + return f"stream-{len(forks)}" + + monkeypatch.setattr(worker, "_m3_fc_fork", fork) + return worker, forks + + +@pytest.mark.parametrize("nreq", [1, 2]) +def test_the_fc_is_forked_at_the_last_captured_layer(monkeypatch, nreq) -> None: + """A fused-verify step of at most 8 verify rows forks the drafter's fc once, right after the last captured layer + (every captured layer still captures), and the fork is left for the verify to join.""" + worker, forks = _fc_early_worker(monkeypatch) + md = _CaptureMetadata(nreq) + drafter = object() + worker.arm_capture_norm(SimpleNamespace(num_seqs=nreq), md, object(), drafter) + for layer_id in md.layers_to_capture[:-1]: + worker.capture_hidden_states(md, layer_id, None) + assert not forks + worker.capture_hidden_states(md, md.layers_to_capture[-1], None) + worker.capture_hidden_states(md, spec_worker.FINAL_HIDDEN_LAYER_ID, None) + assert md.captured == [1, 29, 56, spec_worker.FINAL_HIDDEN_LAYER_ID] + assert forks == [(4 * nreq, drafter, None)] + assert worker._m3_fc_early_stream == "stream-1" and worker._m3_fc_early_drafter is None + + +def test_no_early_fork_with_a_final_state_capture(monkeypatch) -> None: + """The final-state column is written after the model returns, so captures that include it keep the fc at its + stock place instead of forking at the last captured decoder layer.""" + worker, forks = _fc_early_worker(monkeypatch) + md = _CaptureMetadata(1) + md.layers_to_capture = (spec_worker.FINAL_HIDDEN_LAYER_ID, 29, 56) + worker.arm_capture_norm(SimpleNamespace(num_seqs=1), md, object(), object()) + for layer_id in (29, 56, spec_worker.FINAL_HIDDEN_LAYER_ID): + worker.capture_hidden_states(md, layer_id, None) + assert not forks and worker._m3_fc_early_stream is None + + +@pytest.mark.parametrize("nreq,eligible", [(3, True), (8, True), (1, False)]) +def test_no_early_fork_past_its_rows_or_off_the_fused_verify(monkeypatch, nreq, eligible) -> None: + worker, forks = _fc_early_worker(monkeypatch, eligible) + md = _CaptureMetadata(nreq) + worker._m3_fc_early_stream = "stale" + worker.arm_capture_norm(SimpleNamespace(num_seqs=nreq), md, object(), object()) + for layer_id in md.layers_to_capture: + worker.capture_hidden_states(md, layer_id, None) + assert not forks and worker._m3_fc_early_stream is None + + +# ------------------------------------------------------------------ the NVFP4 producer's coverage + + +def test_the_producer_covers_the_tp4_nvfp4_layers_only() -> None: + """``dispatch.PRODUCER`` covers a sparse layer on an NVFP4 cache at the TP4 heads; elsewhere the built-in + producers run without a miss.""" + covered = layers.producer_pw_covered + assert covered(True, 16, 1, 1) + assert not covered(False, 16, 1, 1) + assert not covered(True, 8, 1, 1) + assert not covered(True, 16, 2, 2) + assert dispatch.PRODUCER.covers("nvfp4_tp4") + + +class _FakeMsaAttn: + indexer_kv_dtype = "fp8" + + +def _producer_layer(monkeypatch, nvfp4: bool = True, heads: int = 16): + """What the producer dispatch reads before its first precondition: a fused-projection MSA sparse layer.""" + monkeypatch.setattr(layers, "MiniMaxM3MsaSparseAttention", _FakeMsaAttn) + return SimpleNamespace( + enable_fused_qkv_index_projection=True, + attn=_FakeMsaAttn(), + main_kv_is_fp8=not nvfp4, + main_kv_is_nvfp4=nvfp4, + num_heads=heads, + num_key_value_heads=1, + sparse_num_index_heads=1, + pos_embd_params=None, + ) + + +def _produce(layer): + packed = torch.zeros(4, 20 * 128, dtype=torch.bfloat16) + return layers.M3Attention._fused_fp8_qkv_indexer_norm_rope_kv_insert( + layer, packed, None, SimpleNamespace() + ) + + +def test_a_covered_producer_layer_missing_a_precondition_raises_under_require(monkeypatch) -> None: + """A TP4 NVFP4 sparse layer whose step lacks a producer precondition (here the position ids) is a miss.""" + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=r"'m3_producer' covers this step at 4 rows" + ): + _produce(_producer_layer(monkeypatch)) + + +def test_a_covered_producer_layer_without_the_fused_projection_raises_under_require( + monkeypatch, +) -> None: + layer = _producer_layer(monkeypatch) + layer.enable_fused_qkv_index_projection = False + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="fuse_qkv_index_projection is disabled"): + layers.M3Attention._check_producer_projection(layer) + + +def test_an_uncovered_layer_without_the_fused_projection_is_not_a_miss(monkeypatch) -> None: + layer = _producer_layer(monkeypatch, nvfp4=False) + layer.enable_fused_qkv_index_projection = False + dispatch.set_require(True) + layers.M3Attention._check_producer_projection(layer) + + +def test_a_covered_producer_miss_falls_back_and_logs_once(monkeypatch, log) -> None: + layer = _producer_layer(monkeypatch) + assert _produce(layer) is None and _produce(layer) is None + assert len(log.warning) == 1 and "m3_producer at 4 rows" in log.warning[0] + + +@pytest.mark.parametrize("nvfp4,heads", [(False, 16), (True, 8)], ids=["fp8_cache", "tp8_heads"]) +def test_an_uncovered_producer_layer_falls_back_silently(monkeypatch, log, nvfp4, heads) -> None: + dispatch.set_require(True) + assert _produce(_producer_layer(monkeypatch, nvfp4, heads)) is None + assert not log.warning + + +def test_a_covered_producer_layer_with_the_wrong_pool_layout_is_not_supported() -> None: + """The pool-layout check the covered branch takes before the native fallback (then a miss).""" + pool = torch.zeros(4, 2, 1, 128, 32, dtype=torch.uint8) + scales = torch.zeros(4, 2, 1, 128, 8, dtype=torch.uint8) + assert not layers.producer_pw_supported(16, 1, 1, pool, scales) + good = torch.zeros(4, 2, 1, 128, 64, dtype=torch.uint8) + assert layers.producer_pw_supported(16, 1, 1, good, scales) + assert not layers.producer_pw_supported( + 16, 1, 1, good, scales.transpose(3, 4).contiguous().transpose(3, 4) + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_prebuild.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_prebuild.py new file mode 100644 index 000000000000..34bf7bafe4c2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_prebuild.py @@ -0,0 +1,316 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode kernels ``build_targets`` 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, and the index scorer with and without its early-issue schedule at every decode +query length. + +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 +from unittest import mock + +import pytest +import torch +from torch import nn + +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_103__tp4 import ( # noqa: E501 + dispatch, + layers, + linears, + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__tp4.ops.gemm.m3_mx_gemv import ( # noqa: E501 + O_PROJ, + SHARED_DOWN, + SHARED_GATE_UP, +) +from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.mx_gemv import op as mx_gemv + +assert torch.cuda.is_available(), "the prebuild checks require a CUDA device" + +INDEX_HEADS = 4 +INDEX_DIM = 128 +DRAFT_LEN = 3 + + +@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 _moe(layer_idx: int, shared: bool = True) -> layers.M3MoE: + """An M3MoE built by its own constructor over a stub stock block: only what prebuild_engines and the lookups + read (the builds are stubbed).""" + + def stock_init(self, model_config, aux_stream_dict, layer_idx=None): + nn.Module.__init__(self) + # One rank: the build agreement takes no collective. + self.mapping = SimpleNamespace(tp_size=1, enable_attention_dp=False) + self.shared_experts = None + + with mock.patch.object(layers.MiniMaxM3MoE, "__init__", stock_init): + moe = layers.M3MoE(None, {}, layer_idx) + moe.shared_experts = object() if shared else None + return moe + + +def _model(monkeypatch, num_layers: int = 3, indexer_kv_dtype: str = "fp8"): + """The parts of MiniMaxM3Tp4ModelingV2 prebuild_kernels reads: layers (MoE + sparse attention), the embedding + device and the speculative config.""" + monkeypatch.setattr(layers, "MiniMaxM3MsaSparseAttention", _FakeSparseAttn) + model_layers = [] + for idx in range(num_layers): + model_layers.append( + SimpleNamespace( + block_sparse_moe=_moe(idx, shared=idx != 1) if idx else None, + self_attn=SimpleNamespace( + is_sparse_attention_layer=idx != 0, attn=_FakeSparseAttn(indexer_kv_dtype) + ), + ) + ) + stub = SimpleNamespace( + layers=model_layers, + embed_tokens=SimpleNamespace(weight=torch.empty(1, device="cuda")), + model_config=SimpleNamespace( + spec_config=SimpleNamespace(max_draft_len=DRAFT_LEN, max_total_draft_tokens=DRAFT_LEN) + ), + ) + stub.decode_query_lens = lambda: modeling.MiniMaxM3Tp4ModelingV2.decode_query_lens(stub) + return stub + + +def _stub_engine_builds(monkeypatch, built: list) -> None: + def build(self, num_tokens, shared=True): + built.append(("moe", self.layer_idx, num_tokens, shared)) + return ("engine", self.layer_idx, num_tokens) + + def build16(self, num_tokens): + built.append(("moe16", self.layer_idx, num_tokens)) + return ("engine16", self.layer_idx, num_tokens) + + monkeypatch.setattr(layers.M3MoE, "_build_engine", build) + monkeypatch.setattr(layers.M3MoE, "_build_moe16_engine", build16) + + +def _forbid_compiles(monkeypatch) -> 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(layers.M3MoE, "_build_engine", boom) + monkeypatch.setattr(layers.M3MoE, "_build_moe16_engine", boom) + monkeypatch.setattr(layers._scorer_runner(), "_compile", boom) + + +def test_decode_query_lens_follow_the_speculative_config(monkeypatch) -> None: + model = _model(monkeypatch) + assert model.decode_query_lens() == [1, 1 + DRAFT_LEN] + model.model_config.spec_config = None + assert model.decode_query_lens() == [1] + model.model_config.spec_config = SimpleNamespace( + max_draft_len=4, max_total_draft_tokens=4, draft_len_schedule={1: 4, 8: 2} + ) + assert model.decode_query_lens() == [1, 3, 5] + + +def test_prebuild_builds_every_engine_count_and_scorer_variant(monkeypatch) -> None: + built = [] + _stub_engine_builds(monkeypatch, built) + monkeypatch.setattr(modeling, "adopt_m3_linears", lambda root: built.append(("linears",))) + model = _model(monkeypatch) + times = modeling.MiniMaxM3Tp4ModelingV2.prebuild_kernels(model) + + assert built[0] == ("linears",) + assert mx_gemv.frozen(), "prebuild must close the GEMV compiles" + # Layer 0 is dense, layer 1 has no shared expert (no engine takes it); layer 2 gets every count, in order. + assert built[1:] == [ + ("moe", 2, n, True) for n in range(1, dispatch.MOE_ENGINE_MAX_TOKENS + 1) + ] + [("moe16", 2, n) for n in range(dispatch.MOE16_MIN_TOKENS, dispatch.MOE16_MAX_TOKENS + 1)] + moe = model.layers[2].block_sparse_moe + assert set(moe._engines) == {(n, True, True) for n in dispatch.MOE_ENGINE.keys} + assert set(moe._moe16_engines) == set(dispatch.MOE16_ENGINE.keys) + + runner = layers._scorer_runner() + want = { + (torch.float8_e4m3fn, INDEX_HEADS, qlen, INDEX_DIM, layers.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(runner.kernel_cache) + for key in ("gemv_s", "moe_s", "scorer_s"): + assert times[key] >= 0.0 + + # A second build_targets builds nothing again. + built.clear() + modeling.MiniMaxM3Tp4ModelingV2.prebuild_kernels(model) + assert built == [("linears",)] + + +def test_engine_lookups_after_prebuild_never_build(monkeypatch) -> None: + built = [] + _stub_engine_builds(monkeypatch, built) + moe = _moe(5) + moe.prebuild_engines() + _forbid_compiles(monkeypatch) + for require in (False, True): + dispatch.set_require(require) + for n in sorted(dispatch.MOE_ENGINE.keys): + assert moe.engine(n, shared=True) == ("engine", 5, n) + for n in sorted(dispatch.MOE16_ENGINE.keys): + assert moe.moe16_engine(n) == ("engine16", 5, n) + + +def test_missing_engine_is_a_miss_not_a_build(monkeypatch) -> None: + _forbid_compiles(monkeypatch) + moe = _moe(7) + dispatch.set_require(False) + assert moe.engine(3, shared=True) is None + assert moe.moe16_engine(10) is None + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="not built before warmup"): + moe.engine(3, shared=True) + with pytest.raises(dispatch.FastPathUnavailable, match="not built before warmup"): + moe.moe16_engine(10) + + +def test_failed_engine_build_is_logged_then_missed(monkeypatch) -> None: + def fail(self, num_tokens, shared=True): + raise ValueError("experts backend is not TRTLLM-Gen NVFP4 SwiGLU-OAI") + + monkeypatch.setattr(layers.M3MoE, "_build_engine", fail) + monkeypatch.setattr(layers.M3MoE, "_build_moe16_engine", lambda self, n: fail(self, n)) + moe = _moe(9) + moe.prebuild_engines() + assert all(v is None for v in moe._engines.values()) and len(moe._engines) == 8 + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="build failed"): + moe.engine(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, m3_push=push, tp_size=4) + + +@pytest.mark.parametrize( + "nk,push", + ((O_PROJ, linears.PUSH_ATTN), (SHARED_DOWN, linears.PUSH_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 = linears.gemv_variants(module) + 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 == linears.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) + assert linears.prebuild_m3_gemvs(module, use_pdl) > 0 + 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: + model = _model(monkeypatch, indexer_kv_dtype=indexer_kv_dtype) + keys = layers.prebuild_index_scorer( + model.layers, model.decode_query_lens(), torch.device("cuda") + ) + assert len(keys) == 4 + _forbid_compiles(monkeypatch) + 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() + + +def test_moe_keeps_its_layer_index(): + # prebuild_engines names the layer in its build and agreement messages. + assert _moe(5).layer_idx == 5 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_rank_agreement.py new file mode 100644 index 000000000000..1a906d31f9d3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/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 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 torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4 import ( # noqa: E501 + dispatch, + layers, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 + m3_mc_allreduce as mc, +) + + +@pytest.fixture(autouse=True) +def _restore_modes(): + require = dispatch.require_mode() + yield + dispatch.set_require(require) + dispatch.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 _moe(layer_idx: int, fail: bool) -> layers.M3MoE: + """An M3MoE on two simulated TP ranks whose engine builds are stubbed (and fail on one).""" + + def stock_init(self, model_config, aux_stream_dict, layer_idx=None): + nn.Module.__init__(self) + self.mapping = _mapping(2) + self.shared_experts = None + + with mock.patch.object(layers.MiniMaxM3MoE, "__init__", stock_init): + moe = layers.M3MoE(None, {}, layer_idx) + moe.shared_experts = object() + + def build(num_tokens, shared=True): + if fail: + raise RuntimeError("compile failed on this rank") + return ("engine", num_tokens) + + moe._build_engine = build + moe._build_moe16_engine = lambda num_tokens: build(num_tokens) + return moe + + +@pytest.mark.parametrize("fail", [(False, False), (False, True)]) +def test_engine_build_agreed_across_ranks(two_rank_comm, fail): + moes = [_moe(3, f) for f in fail] + _run_ranks([m.prebuild_engines for m in moes]) + keys = sorted(moes[0]._engines) + sorted(moes[0]._moe16_engines) + assert keys and two_rank_comm.calls == 2 * len(keys) + for engines in ("_engines", "_moe16_engines"): + a, b = getattr(moes[0], engines), getattr(moes[1], engines) + 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 built-in MoE. + dispatch.set_require(False) + for n in sorted(dispatch.MOE_ENGINE.keys): + got = [m.engine(n, shared=True) for m in moes] + 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): + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.spec_worker import ( # noqa: E501 + 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/test_modeling_v2_routing.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py index 90532e2b552d..04896a8f7a3d 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -31,7 +31,11 @@ from pydantic import ValidationError from transformers import PretrainedConfig -from tensorrt_llm._torch._experimental.modeling_v2._router_index import modeling_v2_resolve +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + ModelingV2Mode, + modeling_v2_resolve, + resolved_mode, +) from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_auto import AutoModelForCausalLM from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig, NGramDecodingConfig, TorchLlmArgs @@ -142,6 +146,14 @@ def test_r1_dep4_matches(mode): assert _qualified(modeling_v2_resolve(config, _llm_args(mode))) == _R1_TARGET +@pytest.mark.parametrize("mode", ["auto", "require"]) +def test_the_target_reads_the_mode_it_was_resolved_under(mode): + """A target raises instead of falling back under require, so it has to know which mode built it.""" + config = _model_config(_gpt_oss_config()) + target = modeling_v2_resolve(config, _llm_args(mode)) + assert resolved_mode(target) is ModelingV2Mode(mode) + + def test_resolve_class_reads_the_decision_off_the_config(): """The loader decides and writes the class on the ModelConfig; the class lookup only has to honour it.""" diff --git a/tests/unittest/_torch/sampler/test_greedy_tokens.py b/tests/unittest/_torch/sampler/test_greedy_tokens.py new file mode 100644 index 000000000000..e5100456aadc --- /dev/null +++ b/tests/unittest/_torch/sampler/test_greedy_tokens.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TorchSampler consuming greedy tokens a forward sampled in place of its LM head.""" + +from typing import Any, Callable + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.handle_logits import HandleLogits +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.sampler import TorchSampler +from tensorrt_llm._torch.pyexecutor.sampler.sampler import GREEDY_TOKENS_KEY, greedy_tokens_usable +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.bindings import SamplingConfig +from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams + +VOCAB_SIZE = 16 + + +def _request( + request_id: int = 0, + seq_slot: int = 0, + sampling_params: SamplingParams | None = None, + **kwargs: Any, +) -> LlmRequest: + params = sampling_params if sampling_params is not None else SamplingParams() + return LlmRequest( + request_id=request_id, + max_new_tokens=8, + input_tokens=[1, 2, 3], + sampling_config=SamplingConfig(params._get_sampling_config()), + seq_slot=seq_slot, + is_streaming=False, + **kwargs, + ) + + +def _set(name: str, value: Any) -> Callable[[LlmRequest], None]: + def apply(request: LlmRequest) -> None: + setattr(request, name, value) + + return apply + + +class TestGreedyTokensUsable: + """Whether a generation request is served by the forward's greedy token alone.""" + + pytestmark = pytest.mark.cpu_only + + def test_plain_greedy_request_is_usable(self): + assert greedy_tokens_usable(_request()) + + @pytest.mark.parametrize( + ("kwargs", "mutate"), + [ + pytest.param({}, _set("py_draft_tokens", [5, 6]), id="draft_tokens"), + pytest.param( + {}, _set("py_embedding_bias", torch.zeros(VOCAB_SIZE)), id="embedding_bias" + ), + pytest.param({}, _set("py_bad_words", [[4]]), id="bad_words"), + pytest.param({}, _set("py_no_repeat_ngram_size", 2), id="no_repeat_ngram"), + pytest.param( + {}, _set("py_logits_post_processors", [print]), id="logits_post_processor" + ), + pytest.param({"return_log_probs": True}, None, id="log_probs"), + pytest.param({"return_generation_logits": True}, None, id="generation_logits"), + pytest.param({"stop_words_list": [[13, 14]]}, None, id="multi_token_stop_word"), + ], + ) + def test_a_logits_reader_makes_the_request_unusable(self, kwargs, mutate): + request = _request(**kwargs) + if mutate is not None: + mutate(request) + + assert not greedy_tokens_usable(request) + + @pytest.mark.parametrize( + "sampling_params", + [ + pytest.param(SamplingParams(temperature=0.8, top_k=50), id="top_k"), + pytest.param(SamplingParams(temperature=0.7, top_p=0.9), id="top_p"), + pytest.param(SamplingParams(repetition_penalty=1.3), id="repetition_penalty"), + pytest.param(SamplingParams(presence_penalty=0.5), id="presence_penalty"), + pytest.param(SamplingParams(frequency_penalty=0.5), id="frequency_penalty"), + pytest.param(SamplingParams(min_tokens=2), id="min_length"), + ], + ) + def test_non_greedy_or_penalized_sampling_is_unusable(self, sampling_params): + assert not greedy_tokens_usable(_request(sampling_params=sampling_params)) + + def test_single_token_stop_words_stay_usable(self): + """Single-token stop words are checked on the host, like the end id.""" + assert greedy_tokens_usable(_request(stop_words_list=[[13], [17]])) + + def test_guided_decoding_makes_the_request_unusable(self): + """Guided decoding masks the logits after the forward, so a greedy token cannot serve it.""" + request = _request(guided_decoding_params=GuidedDecodingParams(regex="[0-9]+")) + assert request.py_guided_decoding_params is not None + + assert not greedy_tokens_usable(request) + + @pytest.mark.parametrize("dummy_flag", ["is_attention_dp_dummy", "is_cuda_graph_dummy"]) + def test_a_dummy_request_is_never_usable(self, dummy_flag): + """The stable greedy path takes no batch with a dummy, so such a batch needs logits.""" + request = _request() + setattr(request, dummy_flag, True) + + assert not greedy_tokens_usable(request) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_handle_logits_takes_greedy_only_outputs() -> None: + """The executor hands HandleLogits the outputs' logits before sampling; a greedy-only batch (no request asks for + logits) has none, and HandleLogits must return without reading them.""" + HandleLogits()([], [_request()], None, 1, [0], True) + + +class TestStableGreedyFromGreedyTokens: + """The stable greedy path takes the forward's tokens instead of an argmax over logits.""" + + @staticmethod + def _sampler(max_num_sequences: int = 3) -> TorchSampler: + return TorchSampler( + TorchSampler.Args( + max_seq_len=16, + max_draft_len=0, + max_num_sequences=max_num_sequences, + max_beam_width=1, + max_total_draft_tokens=0, + disable_overlap_scheduler=True, + ) + ) + + @staticmethod + def _batch(*requests: LlmRequest) -> ScheduledRequests: + scheduled = ScheduledRequests() + scheduled.generation_requests = list(requests) + return scheduled + + @staticmethod + def _one_hot_logits(tokens: list[int]) -> torch.Tensor: + logits = torch.zeros((len(tokens), VOCAB_SIZE), device="cuda") + logits[torch.arange(len(tokens)), torch.tensor(tokens)] = 10.0 + return logits + + def test_greedy_tokens_match_the_logits_argmax(self): + tokens = [7, 3] + slots = [2, 0] + results = {} + for name, outputs in ( + ("logits", {"logits": self._one_hot_logits(tokens)}), + ( + "greedy_tokens", + { + # No logits key, as the engine returns greedy-only outputs. + # A padded graph replays more rows than the batch has requests. + GREEDY_TOKENS_KEY: torch.tensor(tokens + [9], dtype=torch.int64, device="cuda"), + }, + ), + ): + sampler = self._sampler() + sampler.store.new_tokens.zero_() + batch = self._batch(*(_request(i, slot) for i, slot in enumerate(slots))) + + ( + requests, + seq_slots_host, + _, + _, + _, + new_tokens_host, + single_step_greedy, + ) = sampler._process_requests(batch, outputs, sampler.store.new_tokens, [0]) + torch.cuda.synchronize() + + assert single_step_greedy + assert requests == batch.generation_requests + assert seq_slots_host.tolist() == slots + results[name] = ( + new_tokens_host.reshape(-1).tolist(), + sampler.store.new_tokens[0, :, 0].tolist(), + ) + + assert results["greedy_tokens"] == results["logits"] + assert results["greedy_tokens"][0] == tokens + assert [results["greedy_tokens"][1][slot] for slot in slots] == tokens + + def test_missing_logits_for_a_batch_that_needs_them_raises(self): + sampler = self._sampler() + request = _request() + request.py_no_repeat_ngram_size = 2 + outputs = {GREEDY_TOKENS_KEY: torch.tensor([4], dtype=torch.int64, device="cuda")} + + with pytest.raises(RuntimeError, match="greedy tokens without logits"): + sampler._process_requests(self._batch(request), outputs, sampler.store.new_tokens, [0]) + + def test_a_batch_padded_with_an_attention_dp_dummy_asks_for_logits(self): + """A dummy makes the batch unusable, so the engine computes logits rather than keep greedy tokens.""" + live = _request(0, 0) + dummy = _request(1, 1) + dummy.is_attention_dp_dummy = True + batch = self._batch(live, dummy) + + assert not all(greedy_tokens_usable(request) for request in batch.generation_requests) diff --git a/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py b/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py new file mode 100644 index 000000000000..1ea8cdaa58f4 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py @@ -0,0 +1,254 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SpecSampler's per-slot store with and without a model's SpecStepFold.""" + +from typing import Optional + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.runtime_hooks import SpecStepFold +from tensorrt_llm._torch.pyexecutor.sampler import TorchSampler +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSampler +from tensorrt_llm.bindings import SamplingConfig + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + +MAX_DRAFT_LEN = 3 +MAX_NUM_SEQUENCES = 6 +SEQ_SLOTS = [4, 1] + + +class _RecordingStepFold: + """A SpecStepFold whose store decisions are fixed per test; it records what the sampler asks.""" + + def __init__(self, rows_stored: bool = False, scatter_taken: bool = False) -> None: + self._rows_stored = rows_stored + self._scatter_taken = scatter_taken + self.store: Optional[object] = None + self.rows_stored_calls: list[tuple[int, list[int]]] = [] + self.scatter_calls: list[dict[str, torch.Tensor]] = [] + + 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: + raise AssertionError("the sampler never binds engine buffers") + + def bind_sampler_store(self, store: object) -> None: + self.store = store + + def begin_step(self, seq_slots: Optional[list[int]], num_rows: int) -> bool: + raise AssertionError("the sampler never begins a step") + + def gather(self, *args: object) -> bool: + raise AssertionError("the sampler never gathers") + + def sampler_rows_stored(self, outputs: dict, num_skip: int, seq_slots: list[int]) -> bool: + self.rows_stored_calls.append((num_skip, list(seq_slots))) + return self._rows_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: + self.scatter_calls.append( + dict( + seq_slots=seq_slots.clone(), + new_tokens=new_tokens.clone(), + next_new_tokens=next_new_tokens.clone(), + new_tokens_lens=new_tokens_lens.clone(), + next_draft_tokens=next_draft_tokens.clone(), + ) + ) + return self._scatter_taken + + +def _fold(**kwargs: bool) -> SpecStepFold: + return _RecordingStepFold(**kwargs) + + +def _sampler() -> SpecSampler: + sampler = SpecSampler( + TorchSampler.Args( + max_seq_len=64, + max_draft_len=MAX_DRAFT_LEN, + max_num_sequences=MAX_NUM_SEQUENCES, + max_beam_width=1, + max_total_draft_tokens=MAX_DRAFT_LEN, + disable_overlap_scheduler=True, + ) + ) + # The store is allocated uninitialized; give it a recognizable baseline. + for index, tensor in enumerate( + ( + sampler.store.new_tokens, + sampler.store.next_new_tokens, + sampler.store.next_draft_tokens, + sampler.store.new_tokens_lens, + ) + ): + tensor.fill_(-100 - index) + return sampler + + +def _store_snapshot(sampler: SpecSampler) -> dict[str, torch.Tensor]: + store = sampler.store + return dict( + new_tokens=store.new_tokens.clone(), + next_new_tokens=store.next_new_tokens.clone(), + next_draft_tokens=store.next_draft_tokens.clone(), + new_tokens_lens=store.new_tokens_lens.clone(), + ) + + +def _assert_store_equals(sampler: SpecSampler, expected: dict[str, torch.Tensor]) -> None: + torch.cuda.synchronize() + for name, tensor in _store_snapshot(sampler).items(): + assert torch.equal(tensor, expected[name]), name + + +def _step(draft_width: int): + requests = [] + for request_id, seq_slot in enumerate(SEQ_SLOTS): + request = LlmRequest( + request_id=request_id, + max_new_tokens=8, + input_tokens=[1, 2, 3], + sampling_config=SamplingConfig(), + seq_slot=seq_slot, + is_streaming=False, + ) + request.py_draft_tokens = [7] * MAX_DRAFT_LEN + requests.append(request) + scheduled = ScheduledRequests() + scheduled.generation_requests = requests + rows = len(requests) + 1 # plus a CUDA-graph padding row the sampler must not store + outputs = dict( + new_tokens=torch.arange(rows * (draft_width + 1), dtype=torch.int32, device="cuda").view( + rows, draft_width + 1 + ), + new_tokens_lens=torch.arange(1, rows + 1, dtype=torch.int32, device="cuda"), + next_draft_tokens=( + 100 + torch.arange(rows * draft_width, dtype=torch.int32, device="cuda") + ).view(rows, draft_width), + next_new_tokens=( + 200 + torch.arange(rows * (draft_width + 1), dtype=torch.int32, device="cuda") + ).view(rows, draft_width + 1), + ) + return scheduled, outputs + + +def _reference_store(before: dict[str, torch.Tensor], outputs: dict) -> dict[str, torch.Tensor]: + """What the sampler has always stored: each live row's outputs at its slot, zero-padded.""" + expected = {name: tensor.clone() for name, tensor in before.items()} + count = len(SEQ_SLOTS) + slots = torch.tensor(SEQ_SLOTS, dtype=torch.long, device="cuda") + + def padded(tensor: torch.Tensor, width: int) -> torch.Tensor: + tensor = tensor[:count, :width] + return torch.nn.functional.pad(tensor, (0, width - tensor.shape[1])) + + accepted_width = expected["new_tokens"].shape[0] + wire_width = expected["next_new_tokens"].shape[0] + draft_width = expected["next_draft_tokens"].shape[1] + expected["new_tokens"].squeeze(-1).T.index_copy_( + 0, slots, padded(outputs["new_tokens"], accepted_width) + ) + expected["next_new_tokens"].squeeze(-1).T.index_copy_( + 0, slots, padded(outputs["next_new_tokens"], wire_width) + ) + expected["new_tokens_lens"].index_copy_(0, slots, outputs["new_tokens_lens"][:count]) + expected["next_draft_tokens"].index_copy_( + 0, slots, padded(outputs["next_draft_tokens"], draft_width) + ) + return expected + + +@pytest.mark.parametrize("draft_width", [MAX_DRAFT_LEN, MAX_DRAFT_LEN - 1]) +def test_without_a_fold_the_store_is_written_as_before(draft_width): + sampler = _sampler() + scheduled, outputs = _step(draft_width) + before = _store_snapshot(sampler) + + state = sampler.sample_async(scheduled, outputs, []) + + assert sampler.spec_step_fold is None + expected = _reference_store(before, outputs) + _assert_store_equals(sampler, expected) + state.sampler_event.synchronize() + assert torch.equal(state.host.new_tokens, expected["new_tokens"].cpu()) + assert torch.equal(state.host.next_draft_tokens, expected["next_draft_tokens"].cpu()) + + +def test_binding_a_fold_hands_it_the_sampler_store(): + sampler = _sampler() + fold = _fold() + + sampler.bind_spec_step_fold(fold) + + assert sampler.spec_step_fold is fold + assert fold.store is sampler.store + + +def test_rows_the_forward_already_stored_are_left_alone(): + sampler = _sampler() + fold = _fold(rows_stored=True) + sampler.bind_spec_step_fold(fold) + scheduled, outputs = _step(MAX_DRAFT_LEN) + before = _store_snapshot(sampler) + + state = sampler.sample_async(scheduled, outputs, []) + + assert fold.rows_stored_calls == [(0, SEQ_SLOTS)] + assert fold.scatter_calls == [] + _assert_store_equals(sampler, before) + # The host copies still read the (fold-owned) store. + state.sampler_event.synchronize() + assert torch.equal(state.host.new_tokens, before["new_tokens"].cpu()) + + +def test_a_fold_scatter_replaces_the_index_copies(): + sampler = _sampler() + fold = _fold(scatter_taken=True) + sampler.bind_spec_step_fold(fold) + scheduled, outputs = _step(MAX_DRAFT_LEN) + before = _store_snapshot(sampler) + + sampler.sample_async(scheduled, outputs, []) + + _assert_store_equals(sampler, before) + assert len(fold.scatter_calls) == 1 + call = fold.scatter_calls[0] + count = len(SEQ_SLOTS) + assert call["seq_slots"].tolist() == SEQ_SLOTS + assert call["seq_slots"].dtype == torch.long and call["seq_slots"].is_cuda + assert torch.equal(call["new_tokens"], outputs["new_tokens"][:count]) + assert torch.equal(call["next_new_tokens"], outputs["next_new_tokens"][:count]) + assert torch.equal(call["new_tokens_lens"], outputs["new_tokens_lens"][:count]) + assert torch.equal(call["next_draft_tokens"], outputs["next_draft_tokens"][:count]) + + +@pytest.mark.parametrize("draft_width", [MAX_DRAFT_LEN, MAX_DRAFT_LEN - 1]) +def test_a_fold_that_declines_leaves_the_sampler_store_path(draft_width): + sampler = _sampler() + fold = _fold(rows_stored=False, scatter_taken=False) + sampler.bind_spec_step_fold(fold) + scheduled, outputs = _step(draft_width) + before = _store_snapshot(sampler) + + sampler.sample_async(scheduled, outputs, []) + + assert fold.rows_stored_calls == [(0, SEQ_SLOTS)] + assert len(fold.scatter_calls) == 1 + _assert_store_equals(sampler, _reference_store(before, outputs)) From c76faa4bd5692aca2e299468f16a010971aadbd7 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:47:59 -0700 Subject: [PATCH 02/14] [None][fix] Satisfy mypy on the MiniMax-M3 sampler fast-greedy path Pick the row source explicitly, raise when neither logits nor greedy tokens are present, and assert logits before the fast greedy kernel. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/sampler/sampler.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 2e2517c259a2..c7b6a5df052e 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -2542,9 +2542,14 @@ def _process_requests( bool, ]: # A forward that sampled greedy tokens in place of the logits (GREEDY_TOKENS_KEY) may return no logits. - raw_logits_cuda = model_outputs.get("logits") - greedy_tokens = model_outputs.get(GREEDY_TOKENS_KEY) - rows_cuda = raw_logits_cuda if greedy_tokens is None else greedy_tokens + raw_logits_cuda: Optional[torch.Tensor] = model_outputs.get("logits") + greedy_tokens: Optional[torch.Tensor] = model_outputs.get(GREEDY_TOKENS_KEY) + if greedy_tokens is not None: + rows_cuda = greedy_tokens + elif raw_logits_cuda is not None: + rows_cuda = raw_logits_cuda + else: + raise RuntimeError("the forward returned neither logits nor greedy tokens") generation_requests = scheduled_requests.generation_requests request_ids = [request.py_request_id for request in generation_requests] @@ -2700,6 +2705,8 @@ def _process_requests( next_tokens, new_tokens_cuda, seq_slots_cuda, self.max_beam_width ) else: + # Without greedy tokens, rows_cuda is the logits. + assert raw_logits_cuda is not None next_tokens = fast_greedy_sample_kernel( raw_logits_cuda[: len(generation_requests)], new_tokens_cuda, From 24cccafc5d7f79d1fb43263245f76910fd2374dd Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 14 Oct 2026 11:56:07 -0700 Subject: [PATCH 03/14] [None][feat] MiniMax-M3 NVFP4 TP8 decode target for modeling_v2 Adds the modeling_v2 target minimax_m3_nvfp4__sm_103__tp8 (GB300, 8 ranks over two trays, tensor-parallel attention and experts, one-model Eagle3 or no speculation) and routes world = tp = moe_tp = 8 to it; the tp4 branch of the router is unchanged. - Decode path: the tp4 target's exact decode path at the TP8 geometry, plus the TP8 items: the 1..8-token MoE engine at expert width 384 (logits replicated 32x at <= 4 tokens), an 8-head one-launch long-context sparse decode staged before its grid wait at <= 8 rows, a fused index score + top-16 select for spans of <= 2 requests, the producer releasing its dependents on entry, the long-row select's early release at <= 8 rows, the TP8 MXFP8 GEMV table, fixed-tactic low-M BF16 GEMMs for the drafter linears, the drafter decode's KV split capped at 16K for one request, BF16 target and draft heads, and the drafter fc forked beside the verify head. - Isolation: the target imports nothing from the tp4 target; the MoE engine and the MX GEMV op it changes are its own copies inside the target. The two new kernels are compiled ops (cpp/tensorrt_llm/kernels/minimaxM3Tp8, trtllm::minimax_m3_tp8_*). A drift test pins the files copied unchanged from the tp4 target. - fmha/flashinfer_trtllm_gen.py applies a layer's KV-split cap (msa_decode.dense_decode_max_seq_len); a no-op unless a layer sets one. - Same contract as the tp4 target: every kernel a catalog entry with CELLS and a torch reference, dispatch.py listing each fast path's rows, a covered shape that cannot run raising under modeling_v2: require, kernels built before warmup, every item on by default with one constant to turn it off. - Numerics match stock up to summation order (BF16 heads and drafter, NVFP4 experts, MXFP8 projections). Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../kernels/minimaxM3Tp8/fusedScoreSelect.cu | 770 ++++ .../kernels/minimaxM3Tp8/fusedScoreSelect.h | 77 + .../kernels/minimaxM3Tp8/sparseDecodeLc8.cu | 610 +++ .../kernels/minimaxM3Tp8/sparseDecodeLc8.h | 75 + cpp/tensorrt_llm/thop/CMakeLists.txt | 2 + .../thop/minimaxM3Tp8FusedScoreSelectOp.cpp | 173 + .../thop/minimaxM3Tp8SparseDecodeLc8Op.cpp | 139 + .../minimax_m3_nvfp4__sm_103__tp8/__init__.py | 3 + .../minimax_m3_nvfp4__sm_103__tp8/dispatch.py | 527 +++ .../kernels/__init__.py | 3 + .../kernels/long_select.py | 62 + .../kernels/m3_attn_lc8.py | 110 + .../kernels/m3_draft_lowm.py | 121 + .../kernels/m3_draft_step.py | 150 + .../kernels/m3_fused_score_select.py | 110 + .../kernels/m3_index_score.py | 645 +++ .../kernels/m3_mc_allreduce.py | 368 ++ .../kernels/m3_moe/__init__.py | 16 + .../kernels/m3_moe/m3_moe_kernel.py | 3817 +++++++++++++++++ .../kernels/m3_moe/op.py | 677 +++ .../kernels/m3_prequant.py | 111 + .../kernels/m3_router_gemv.py | 107 + .../kernels/m3_sampled_accept.py | 158 + .../kernels/m3_shard_accept.py | 176 + .../kernels/m3_spec_tail.py | 113 + .../kernels/msa_kv_update.py | 89 + .../kernels/mx_gemv/__init__.py | 4 + .../kernels/mx_gemv/m3_mx_gemv_kernel.py | 1236 ++++++ .../kernels/mx_gemv/op.py | 472 ++ .../kernels/triton_sparse_decode.py | 1662 +++++++ .../kernels/triton_sparse_decode_mq.py | 237 + .../minimax_m3_nvfp4__sm_103__tp8/layers.py | 1403 ++++++ .../minimax_m3_nvfp4__sm_103__tp8/linears.py | 232 + .../minimax_m3_nvfp4__sm_103__tp8/modeling.py | 624 +++ .../ops/__init__.py | 3 + .../ops/attention/__init__.py | 3 + .../ops/attention/m3_dense_decode.py | 187 + .../ops/attention/m3_fused_score_select.py | 189 + .../ops/attention/m3_index_score.py | 153 + .../ops/attention/m3_producer.py | 197 + .../ops/attention/m3_select_long.py | 172 + .../ops/attention/m3_sparse_decode.py | 215 + .../ops/attention/m3_sparse_decode_lc8.py | 433 ++ .../ops/attention/msa_kv_update.py | 200 + .../ops/comm/__init__.py | 3 + .../ops/comm/m3_mc_allreduce_norm.py | 277 ++ .../ops/comm/m3_mc_finalize_allreduce_norm.py | 197 + .../ops/gemm/__init__.py | 3 + .../ops/gemm/m3_draft_lowm.py | 150 + .../ops/gemm/m3_mx_gemv.py | 160 + .../ops/gemm/m3_mx_gemv_bf16.py | 106 + .../ops/moe/__init__.py | 3 + .../ops/moe/m3_moe_engine.py | 325 ++ .../ops/moe/m3_router_gemv.py | 101 + .../ops/moe/m3_router_gemv_fp4.py | 63 + .../ops/moe/m3_router_tc.py | 106 + .../ops/norm/__init__.py | 3 + .../ops/norm/m3_embedding_firstnorm.py | 89 + .../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_step.py | 359 ++ .../ops/spec/m3_sampled_accept.py | 269 ++ .../ops/spec/m3_sampled_accept_ws_words.py | 43 + .../ops/spec/m3_shard_accept.py | 284 ++ .../ops/spec/m3_slot_scatter.py | 72 + .../ops/spec/m3_step_gather.py | 152 + .../minimax_m3_nvfp4__sm_103__tp8/runtime.py | 50 + .../spec_fold.py | 112 + .../spec_worker.py | 1119 +++++ .../minimax_m3_nvfp4__sm_103__tp8/weights.py | 62 + .../modeling_v2/models/minimax_m3/routing.py | 44 +- .../backends/fmha/flashinfer_trtllm_gen.py | 3 + .../attention/_m3_nvfp4_cache.py | 105 + .../test_modeling_v2_m3_dense_decode.py | 152 + .../test_modeling_v2_m3_fused_score_select.py | 230 + .../test_modeling_v2_m3_index_score.py | 94 + .../attention/test_modeling_v2_m3_producer.py | 179 + .../test_modeling_v2_m3_select_long.py | 159 + .../test_modeling_v2_m3_sparse_decode.py | 144 + .../test_modeling_v2_m3_sparse_decode_lc8.py | 195 + .../test_modeling_v2_msa_kv_update.py | 128 + .../comm/_m3_draft_ar_norm_op_matrix.py | 394 ++ .../comm/_m3_mc_allreduce_norm_op_matrix.py | 670 +++ .../comm/_rank_job.py | 107 + .../comm/test_modeling_v2_m3_draft_ar_norm.py | 38 + ...eling_v2_m3_mc_allreduce_norm_op_matrix.py | 60 + .../gemm/test_modeling_v2_m3_draft_lowm.py | 196 + .../gemm/test_modeling_v2_m3_mx_gemv.py | 91 + .../gemm/test_modeling_v2_m3_mx_gemv_bf16.py | 67 + .../moe/_m3_moe_operands.py | 172 + .../moe/test_modeling_v2_m3_moe_engine.py | 134 + .../moe/test_modeling_v2_m3_router_gemv.py | 148 + ...test_modeling_v2_m3_embedding_firstnorm.py | 71 + .../spec/_m3_spec_inputs.py | 106 + .../spec/test_modeling_v2_m3_capture_norm.py | 57 + .../spec/test_modeling_v2_m3_chunk_norm.py | 50 + .../spec/test_modeling_v2_m3_draft_input.py | 98 + .../spec/test_modeling_v2_m3_draft_step.py | 611 +++ .../test_modeling_v2_m3_exact_numerics.py | 305 ++ .../test_modeling_v2_m3_sampled_accept.py | 167 + .../spec/test_modeling_v2_m3_shard_accept.py | 120 + .../spec/test_modeling_v2_m3_slot_scatter.py | 57 + .../spec/test_modeling_v2_m3_step_gather.py | 80 + .../test_modeling_v2_m3_capture_boundary.py | 65 + .../test_modeling_v2_m3_dispatch.py | 995 +++++ .../test_modeling_v2_m3_prebuild.py | 305 ++ .../test_modeling_v2_m3_rank_agreement.py | 199 + .../test_modeling_v2_m3_tp4_drift.py | 121 + .../modeling_v2/test_modeling_v2_routing.py | 18 + 111 files changed, 28153 insertions(+), 15 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Tp8/fusedScoreSelect.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Tp8/fusedScoreSelect.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Tp8/sparseDecodeLc8.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Tp8/sparseDecodeLc8.h create mode 100644 cpp/tensorrt_llm/thop/minimaxM3Tp8FusedScoreSelectOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3Tp8SparseDecodeLc8Op.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/dispatch.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/long_select.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_attn_lc8.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_draft_lowm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_draft_step.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_fused_score_select.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_mc_allreduce.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/m3_moe_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_prequant.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_router_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_sampled_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_shard_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_spec_tail.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/msa_kv_update.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/m3_mx_gemv_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/triton_sparse_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/triton_sparse_decode_mq.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/layers.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/linears.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_dense_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_fused_score_select.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_producer.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_select_long.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode_lc8.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/msa_kv_update.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_mc_allreduce_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp8/ops/gemm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_draft_lowm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv_bf16.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_moe_engine.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_router_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_router_gemv_fp4.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_router_tc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/m3_embedding_firstnorm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_capture_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_chunk_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_draft_input.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_draft_step.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_sampled_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp8/ops/spec/m3_shard_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_slot_scatter.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_step_gather.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/runtime.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec_fold.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec_worker.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/weights.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/_m3_nvfp4_cache.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_dense_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_fused_score_select.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_index_score.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_producer.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_select_long.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_sparse_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp8/attention/test_modeling_v2_msa_kv_update.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_draft_ar_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_rank_job.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_draft_ar_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp8/gemm/test_modeling_v2_m3_draft_lowm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mx_gemv.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp8/moe/_m3_moe_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_moe_engine.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_gemv.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/_m3_spec_inputs.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_capture_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_chunk_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_input.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_step.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_exact_numerics.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_sampled_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_shard_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_slot_scatter.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_step_gather.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_capture_boundary.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_dispatch.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_prebuild.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Tp8/fusedScoreSelect.cu b/cpp/tensorrt_llm/kernels/minimaxM3Tp8/fusedScoreSelect.cu new file mode 100644 index 000000000000..dfc570239a5c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Tp8/fusedScoreSelect.cu @@ -0,0 +1,770 @@ +/* + * 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 fused index score + top-16 block select in one launch. +// +// Scoring: grid (splits, batch), one CTA per SM. CTA (s, b) owns blocks s, s + splits, ... of request b and TMA-loads +// their index-K pages through a stage ring (one producer warp, 8 consumer warps; each consumer warp scores a 32-row +// slice of a page against all q columns with m16n8k16 FP16 MMAs, column = head * 4 + query row). Block scores go to a +// row-major scratch; the last CTA of the request (arrival counter) bulk-copies the request's rows back into its idle +// stage ring and selects the top 16 per (query row, KV head): a threshold bound from per-thread maxima, a candidate +// list, exact ranking by (key desc, block asc), and an exact radix select when ties overflow the candidate list. + +#include "tensorrt_llm/kernels/minimaxM3Tp8/fusedScoreSelect.h" + +#include "tensorrt_llm/common/cudaDriverWrapper.h" + +#include +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr int kTopK = 16; +constexpr int kBlockK = 128; +constexpr int kHeadDim = 128; +constexpr int kStageBytes = kBlockK * kHeadDim; // fp8 +constexpr int kConsumerWarps = 8; +constexpr int kConsumerThreads = kConsumerWarps * 32; +constexpr int kThreads = kConsumerThreads + 32; // + producer warp +constexpr int kMaxDql = 4; // B-operand column = head * kMaxDql + query row +constexpr int kCols = 8; // one n8 tile: heads * kMaxDql <= 8 +constexpr int kMaxRows = 8; // kv_heads * dql +constexpr int kMaxLocal = 128; // blocks per CTA +constexpr int kCandCap = 256; +constexpr int kEpiBytes = kMaxLocal * 4 * kCols * 4; +constexpr float kInitScore = 1.0e30F; +constexpr float kLocalScore = 1.0e29F; +constexpr int kNStage = 12; // TMA stage ring depth (pages in flight per CTA) + +using Params = MinimaxM3Tp8FusedScoreSelectParams; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +__device__ __forceinline__ uint32_t twiddle(float x) +{ + uint32_t const bits = __float_as_uint(x); + return bits ^ ((bits & 0x80000000u) ? 0xFFFFFFFFu : 0x80000000u); +} + +__device__ __forceinline__ uint32_t cvt_e4m3x2(uint32_t pair16) +{ + uint32_t out; + unsigned short const in = static_cast(pair16 & 0xFFFFu); + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(out) : "h"(in)); + return out; +} + +__device__ __forceinline__ void mma16816( + float (&c)[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.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"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void mbar_init(uint32_t bar, int count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(bar), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbar_expect_tx(uint32_t bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(bar), "r"(bytes) : "memory"); +} + +__device__ __forceinline__ void mbar_arrive(uint32_t bar) +{ + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" ::"r"(bar) : "memory"); +} + +__device__ __forceinline__ void mbar_wait(uint32_t bar, uint32_t parity) +{ + uint32_t done; + do + { + asm volatile( + "{\n\t.reg .pred p;\n\tmbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n\tselp.u32 %0, 1, 0, p;\n\t}" + : "=r"(done) + : "r"(bar), "r"(parity) + : "memory"); + } while (!done); +} + +__device__ __forceinline__ void tma_load_page( + uint32_t dst, CUtensorMap const* tmap, int page, uint32_t bar, uint64_t policy) +{ + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint" + " [%0], [%1, {%2, %3, %4}], [%5], %6;" ::"r"(dst), + "l"(reinterpret_cast(tmap)), "r"(0), "r"(0), "r"(page), "r"(bar), "l"(policy) + : "memory"); +} + +// Exact radix select over n smem keys (fallback when the threshold bound leaves too many candidates). +__device__ void radix_emit(uint32_t const* kr, int n, uint32_t* hist, int32_t* dst, uint32_t neg_inf_key) +{ + int const lane = threadIdx.x & 31; + uint32_t prefix = 0, mask = 0; + int need = kTopK; + bool exact_bin = false; + for (int shift = 24; shift >= 0; shift -= 8) + { + for (int i = lane; i < 256; i += 32) + hist[i] = 0; + __syncwarp(); + for (int i = lane; i < n; i += 32) + { + uint32_t const k = kr[i]; + if ((k & mask) == prefix) + atomicAdd(&hist[(k >> shift) & 255u], 1u); + } + __syncwarp(); + uint32_t c[8]; + uint32_t sum = 0; +#pragma unroll + for (int j = 0; j < 8; ++j) + { + c[j] = hist[255 - 8 * lane - j]; + sum += c[j]; + } + uint32_t incl = sum; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) + { + uint32_t const v = __shfl_up_sync(0xFFFFFFFFu, incl, o); + if (lane >= o) + incl += v; + } + uint32_t above = incl - sum; + int found_bin = -1; + uint32_t found_above = 0, found_cnt = 0; +#pragma unroll + for (int j = 0; j < 8; ++j) + { + if (found_bin < 0 && above < static_cast(need) && above + c[j] >= static_cast(need)) + { + found_bin = 255 - 8 * lane - j; + found_above = above; + found_cnt = c[j]; + } + above += c[j]; + } + unsigned const bal = __ballot_sync(0xFFFFFFFFu, found_bin >= 0); + int const src = __ffs(bal) - 1; + int const bin = __shfl_sync(0xFFFFFFFFu, found_bin, src); + uint32_t const fa = __shfl_sync(0xFFFFFFFFu, found_above, src); + uint32_t const fc = __shfl_sync(0xFFFFFFFFu, found_cnt, src); + __syncwarp(); + need -= static_cast(fa); + prefix |= static_cast(bin) << shift; + mask |= 255u << shift; + if (fc == static_cast(need)) + { + exact_bin = true; + break; + } + } + int emitted = 0, ties_used = 0; + uint32_t const lt = (1u << lane) - 1u; + for (int c0 = 0; c0 < n; c0 += 32) + { + int const i = c0 + lane; + bool const valid = i < n; + uint32_t const k = valid ? kr[i] : 0u; + uint32_t const km = k & mask; + bool const neg = valid && k == neg_inf_key; + bool sel = valid && (km > prefix || (exact_bin && km == prefix)); + bool const tie = valid && !exact_bin && km == prefix; + unsigned const tie_bal = __ballot_sync(0xFFFFFFFFu, tie); + if (tie) + sel = ties_used + __popc(tie_bal & lt) < need; + ties_used += __popc(tie_bal); + bool const out_ok = sel && !neg; + unsigned const ob = __ballot_sync(0xFFFFFFFFu, out_ok); + if (out_ok) + dst[emitted + __popc(ob & lt)] = i; + emitted += __popc(ob); + } + if (lane >= emitted && lane < kTopK) + dst[lane] = -1; + __syncwarp(); +} + +// Descending bitonic sort of one value per lane (lane 0 ends with the largest). +__device__ __forceinline__ uint32_t warp_sort_desc(uint32_t v, int lane) +{ +#pragma unroll + for (int k = 2; k <= 32; k <<= 1) + { +#pragma unroll + for (int j = k >> 1; j > 0; j >>= 1) + { + uint32_t const o = __shfl_xor_sync(0xFFFFFFFFu, v, j); + bool const desc_block = (lane & k) == 0; + bool const low = (lane & j) == 0; + v = (low == desc_block) ? max(v, o) : min(v, o); + } + } + return v; +} + +__device__ __forceinline__ uint32_t forced_key(float sc, int i, int init_blocks, int local_start) +{ + if (i < init_blocks) + sc = kInitScore; + if (i >= local_start) + sc = kLocalScore; + return twiddle(sc); +} + +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +template +__global__ void __launch_bounds__(kThreads, 1) + m3_fused_score_select_kernel(__grid_constant__ CUtensorMap const tmap, Params p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ uint8_t smem_raw[]; + uint8_t* smem = smem_raw + ((1024u - (smem_u32(smem_raw) & 1023u)) & 1023u); + uint8_t* stages = smem; + float* epi = reinterpret_cast(smem + NSTAGE * kStageBytes); // [kMaxLocal][4 slices][kCols] + uint64_t* bars + = reinterpret_cast(smem + NSTAGE * kStageBytes + kEpiBytes); // full[NSTAGE], empty[NSTAGE] + __shared__ uint32_t s_wtop[kMaxRows][kConsumerWarps][16]; + __shared__ int s_cnt[kMaxRows]; + __shared__ int32_t s_sel_blk[kMaxRows][kTopK]; + __shared__ int s_sel_neg[kMaxRows][kTopK]; + __shared__ int s_last; + __shared__ int32_t s_sb[kMaxRows][32]; + __shared__ __align__(8) uint64_t s_selbar; + + int const s = blockIdx.x, b = blockIdx.y, splits = gridDim.x; + int const tid = threadIdx.x, warp = tid >> 5, lane = tid & 31; + int const g = lane >> 2, tq = lane & 3; + uint32_t const full0 = smem_u32(bars), empty0 = smem_u32(bars + NSTAGE); + + if (tid == 0) + { + for (int i = 0; i < NSTAGE; ++i) + { + mbar_init(full0 + 8 * i, 1); + mbar_init(empty0 + 8 * i, 4); + } + mbar_init(smem_u32(&s_selbar), 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + if (warp == kConsumerWarps && lane == 0) + asm volatile("prefetch.tensormap [%0];" ::"l"(reinterpret_cast(&tmap)) : "memory"); + __syncthreads(); + + // The producer lanes fetch their first page ids together with seq_lens (bounded by the table width), so a CTA + // pays one dependent round trip before its loads go out. That happens, and the loads of the pages that hold only + // tokens older than this step go out, before the grid dependency wait. + int pid_first = 0; + int seqlen = 0; + uint64_t policy = 0; + if (warp == kConsumerWarps) + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" : "=l"(policy)); + int issued = 0; // producer lane: its first page is already in flight + { + if (warp == kConsumerWarps && lane < NSTAGE) + { + int const blk = s + lane * splits; + if (blk < p.bt_width) + pid_first = __ldg(p.bt + static_cast(b) * p.bt_stride + blk); + } + seqlen = p.seq_lens[b]; + if (warp == kConsumerWarps) + { + int const nb0 = seqlen > 0 ? min((seqlen + kBlockK - 1) / kBlockK, p.max_tiles) : 0; + int const my0 = s < nb0 ? (nb0 - s + splits - 1) / splits : 0; + int const stable = max(seqlen - p.dql, 0) / kBlockK; // blocks holding only older tokens + int const blk = s + lane * splits; + if (lane < min(my0, NSTAGE) && blk < stable) + { + mbar_expect_tx(full0 + 8 * lane, kStageBytes); + tma_load_page(smem_u32(stages + lane * kStageBytes), &tmap, pid_first, full0 + 8 * lane, policy); + issued = 1; + } + } + } + + asm volatile("griddepcontrol.wait;" ::: "memory"); + + // The select's per-row valid-block count, loaded by every CTA alongside seq_lens (only the last one uses it). + int const sel_tpr = p.dql * p.kv_heads <= 4 ? kConsumerThreads / 4 : kConsumerThreads / 8; + int const nv_raw = tid < kConsumerThreads && tid / sel_tpr < p.dql * p.kv_heads + ? p.n_valid[b * p.dql + (tid / sel_tpr) % p.dql] + : 0; + int const nb = seqlen > 0 ? min((seqlen + kBlockK - 1) / kBlockK, p.max_tiles) : 0; + if (nb == 0) + { + // Nothing to score: every row of the request is all -1. + if (s == 0 && tid < p.dql * p.kv_heads * kTopK) + { + int const r = tid / kTopK, j = tid % kTopK; + int const t = b * p.dql + r % p.dql, kvh = r / p.dql; + p.out[static_cast(kvh) * p.out_h + static_cast(t) * p.out_t + j] = -1; + } + return; + } + if (s >= nb) + return; + int const parts = min(splits, nb); + int const my = (nb - s + splits - 1) / splits; + int const q_start = seqlen - p.dql; + int const group = p.heads / p.kv_heads; + int const rows = p.dql * p.kv_heads; + + if (warp == kConsumerWarps) + { + if (lane < min(my, NSTAGE) && !issued) + { + mbar_expect_tx(full0 + 8 * lane, kStageBytes); + tma_load_page(smem_u32(stages + lane * kStageBytes), &tmap, pid_first, full0 + 8 * lane, policy); + } + if (lane == 0) + { + for (int j = NSTAGE; j < my; ++j) + { + int const st = j % NSTAGE; + int const page = __ldg(p.bt + static_cast(b) * p.bt_stride + s + j * splits); + mbar_wait(empty0 + 8 * st, ((j / NSTAGE) - 1) & 1); + mbar_expect_tx(full0 + 8 * st, kStageBytes); + tma_load_page(smem_u32(stages + st * kStageBytes), &tmap, page, full0 + 8 * st, policy); + } + } + } + else + { + // Q fragments: column g is (head = g / kMaxDql, row = g % kMaxDql). + uint32_t qlo[4][2], qhi[4][2]; + { + int const head = g / kMaxDql, qpos = g % kMaxDql; + bool const valid = head < p.heads && qpos < p.dql; + uint8_t const* qrow = p.q + static_cast(b * p.dql + qpos) * p.q_st + head * p.q_sh; + uint32_t w0[4], w1[4]; +#pragma unroll + for (int k = 0; k < 4; ++k) + { + w0[k] = valid ? *reinterpret_cast(qrow + 32 * k + 4 * tq) : 0u; + w1[k] = valid ? *reinterpret_cast(qrow + 32 * k + 16 + 4 * tq) : 0u; + } +#pragma unroll + for (int k = 0; k < 4; ++k) + { + qlo[k][0] = cvt_e4m3x2(w0[k]); + qhi[k][0] = cvt_e4m3x2(w0[k] >> 16); + qlo[k][1] = cvt_e4m3x2(w1[k]); + qhi[k][1] = cvt_e4m3x2(w1[k] >> 16); + } + } + int const slice = warp & 3; +#pragma unroll 1 + for (int lb = warp >> 2; lb < my; lb += kConsumerWarps / 4) + { + int const st = lb % NSTAGE; + mbar_wait(full0 + 8 * st, (lb / NSTAGE) & 1); + uint8_t const* stg = stages + st * kStageBytes; + float acc[2][4]; +#pragma unroll + for (int m = 0; m < 2; ++m) +#pragma unroll + for (int e = 0; e < 4; ++e) + acc[m][e] = 0.f; +#pragma unroll + for (int k = 0; k < 4; ++k) + { +#pragma unroll + for (int m = 0; m < 2; ++m) + { + int const ra = 32 * slice + 16 * m + g; // rb = ra + 8 has the same (row & 7) + int const sw = ra & 7; + uint8_t const* pa = stg + ra * 128; + uint8_t const* pb = pa + 8 * 128; + uint32_t const d0 = *reinterpret_cast(pa + (((2 * k) ^ sw) << 4) + 4 * tq); + uint32_t const d1 = *reinterpret_cast(pb + (((2 * k) ^ sw) << 4) + 4 * tq); + uint32_t const d2 = *reinterpret_cast(pa + (((2 * k + 1) ^ sw) << 4) + 4 * tq); + uint32_t const d3 = *reinterpret_cast(pb + (((2 * k + 1) ^ sw) << 4) + 4 * tq); + uint32_t const l0 = cvt_e4m3x2(d0), l1 = cvt_e4m3x2(d1), l2 = cvt_e4m3x2(d2), l3 = cvt_e4m3x2(d3); + uint32_t const h0 = cvt_e4m3x2(d0 >> 16), h1 = cvt_e4m3x2(d1 >> 16), h2 = cvt_e4m3x2(d2 >> 16), + h3 = cvt_e4m3x2(d3 >> 16); + mma16816(acc[m], l0, l1, l2, l3, qlo[k][0], qlo[k][1]); + mma16816(acc[m], h0, h1, h2, h3, qhi[k][0], qhi[k][1]); + } + } + __syncwarp(); + if (lane == 0) + mbar_arrive(empty0 + 8 * st); + // Causal mask + max over this warp's 32 positions. + int const blk = s + lb * splits; + int const pos0 = blk * kBlockK + 32 * slice; +#pragma unroll + for (int j = 0; j < 2; ++j) + { + int const col = 2 * tq + j; + int const qabs = q_start + col % kMaxDql; + float v = -INFINITY; +#pragma unroll + for (int m = 0; m < 2; ++m) + { + float const a = qabs >= pos0 + 16 * m + g ? acc[m][j] : -INFINITY; + float const c = qabs >= pos0 + 16 * m + g + 8 ? acc[m][2 + j] : -INFINITY; + v = fmaxf(v, fmaxf(a, c)); + } +#pragma unroll + for (int o = 4; o < 32; o <<= 1) + v = fmaxf(v, __shfl_xor_sync(0xFFFFFFFFu, v, o)); + if (g == 0) + epi[(lb * 4 + slice) * kCols + col] = v; + } + } + } + __syncthreads(); + + // Per block: max over the 4 slices, amax over each KV head's index heads, row-major scratch stores. + for (int lb = tid; lb < my; lb += kThreads) + { + int const blk = s + lb * splits; + float v[kCols]; +#pragma unroll + for (int c = 0; c < kCols; ++c) + { + float x = epi[(lb * 4) * kCols + c]; +#pragma unroll + for (int sl = 1; sl < 4; ++sl) + x = fmaxf(x, epi[(lb * 4 + sl) * kCols + c]); + v[c] = x; + } + float* dst = p.scr + b * p.scr_b + blk; +#pragma unroll + for (int r = 0; r < kMaxRows; ++r) + { + if (r < rows) + { + int const kvh = r / p.dql, qpos = r % p.dql; + float x = -INFINITY; +#pragma unroll + for (int c = 0; c < kCols; ++c) + if (c / kMaxDql >= kvh * group && c / kMaxDql < (kvh + 1) * group && c % kMaxDql == qpos) + x = fmaxf(x, v[c]); + dst[r * p.scr_r] = x; + } + } + } + __syncthreads(); + + __syncthreads(); + if (tid == 0) + { + int old; + asm volatile("atom.add.acq_rel.gpu.s32 %0, [%1], 1;" : "=r"(old) : "l"(p.counters + b) : "memory"); + s_last = old == parts - 1; + } + __syncthreads(); + if (!s_last) + return; + asm volatile("griddepcontrol.launch_dependents;"); + + // ---- last CTA of request b: top-16 of every (query row, KV head) ---- + // Row r is owned by tpr threads (2 warps for <= 4 rows, 1 warp for <= 8); thread rt of the row takes blocks + // rt, rt + tpr, ... and stages their forced keys in the (now idle) stage ring. + if (tid == 0) + p.counters[b] = 0; + if (tid < kMaxRows) + s_cnt[tid] = 0; + int const tpr = sel_tpr; + int const r = tid / tpr, rt = tid % tpr; + bool const active = tid < kConsumerThreads && r < rows; + int const ldk = (p.max_tiles + 3) & ~3; + uint32_t* keys = reinterpret_cast(stages) + (active ? r : 0) * ldk; + // One bulk copy per row brings the request's block scores (blocks [0, nb), contiguous per row) into the stage + // ring, so the whole row arrives in a single L2 round trip. + uint32_t const selbar = smem_u32(&s_selbar); + if (tid == 0) + { + asm volatile("fence.proxy.async.global;" ::: "memory"); + uint32_t const bytes = static_cast((nb + 3) & ~3) * 4u; + mbar_expect_tx(selbar, bytes * rows); + for (int rr = 0; rr < rows; ++rr) + { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + smem_u32(stages) + static_cast(rr * ldk) * 4u), + "l"(p.scr + b * p.scr_b + rr * p.scr_r), "r"(bytes), "r"(selbar) + : "memory"); + } + } + int n = 0, lst = 0; + uint32_t mx = 0u; + if (active) + { + n = max(0, min(nv_raw, p.max_tiles)); + lst = max(nv_raw - p.local_blocks, 0); + mbar_wait(selbar, 0); + float const* rowf = reinterpret_cast(keys); +#pragma unroll 1 + for (int i0 = rt; i0 < n; i0 += 8 * tpr) + { + float x[8]; +#pragma unroll + for (int u = 0; u < 8; ++u) + x[u] = i0 + u * tpr < n ? rowf[i0 + u * tpr] : 0.f; +#pragma unroll + for (int u = 0; u < 8; ++u) + { + int const i = i0 + u * tpr; + if (i < n) + { + uint32_t const k = forced_key(x[u], i, p.init_blocks, lst); + keys[i] = k; + mx = max(mx, k); + } + } + } + } + // T0 = 16th largest thread maximum of the row: 16 threads each hold a key >= T0, so every top-16 key (ties + // included) is >= T0 (a row with fewer than 16 keys gets T0 = 0, i.e. all of them). + uint32_t const sorted = warp_sort_desc(mx, lane); + uint32_t t0 = __shfl_sync(0xFFFFFFFFu, sorted, kTopK - 1); + if (tpr == 64 && active && lane < kTopK) + s_wtop[r][rt >> 5][lane] = sorted; + __syncthreads(); + uint32_t* cand_key = reinterpret_cast(epi); // [rows][kCandCap] + int32_t* cand_blk = reinterpret_cast(epi) + kMaxRows * kCandCap; + if (active) + { + if (tpr == 64) + { + // 16th largest of two descending 16-lists A, B: max over i of min(A[i-1], B[15-i]), A[-1] = B[-1] = inf. + uint32_t const* A = s_wtop[r][0]; + uint32_t const* B = s_wtop[r][1]; + uint32_t x = 0u; + if (lane <= kTopK) + x = min(lane > 0 ? A[lane - 1] : 0xFFFFFFFFu, lane < kTopK ? B[kTopK - 1 - lane] : 0xFFFFFFFFu); + t0 = __reduce_max_sync(0xFFFFFFFFu, x); + } +#pragma unroll 1 + for (int i0 = rt; i0 < n; i0 += 8 * tpr) + { + uint32_t kk[8]; +#pragma unroll + for (int u = 0; u < 8; ++u) + kk[u] = i0 + u * tpr < n ? keys[i0 + u * tpr] : 0u; + uint32_t hit = 0; +#pragma unroll + for (int u = 0; u < 8; ++u) + hit |= static_cast(kk[u] >= t0 && kk[u] != 0u) << u; + if (hit) + { +#pragma unroll + for (int u = 0; u < 8; ++u) + { + if ((hit >> u) & 1u) + { + int const at = atomicAdd(&s_cnt[r], 1); + if (at < kCandCap) + { + cand_key[r * kCandCap + at] = kk[u]; + cand_blk[r * kCandCap + at] = i0 + u * tpr; + } + } + } + } + } + } + __syncthreads(); + // One warp per row: exact rank by (key desc, block asc), drop -inf, emit ascending ids with a -1 tail. + if (warp < rows) + { + int const rr = warp; + int const t = b * p.dql + rr % p.dql, kvh = rr / p.dql; + int32_t* dst = p.out + static_cast(kvh) * p.out_h + static_cast(t) * p.out_t; + int const c = s_cnt[rr]; + uint32_t const neg_key = twiddle(-INFINITY); + if (c <= 32) + { + // Lane j ranks candidate j against all of them (smem broadcasts), then places the selected ids by + // counting the smaller selected ids. + uint32_t const* ck = cand_key + rr * kCandCap; + int32_t const* cb = cand_blk + rr * kCandCap; + uint32_t const kj = lane < c ? ck[lane] : 0u; + int const bj = lane < c ? cb[lane] : 0x7FFFFFFF; + int rank = 0; +#pragma unroll 8 + for (int o = 0; o < 32; ++o) + { + uint32_t const ko = ck[o]; + int const bo = cb[o]; + rank += (o < c) & ((ko > kj) | ((ko == kj) & (bo < bj))); + } + bool const sel = lane < c && rank < kTopK && kj != neg_key; + int const sb = sel ? bj : 0x7FFFFFFF; + s_sb[rr][lane] = sb; + __syncwarp(); + int pos = 0; +#pragma unroll 8 + for (int o = 0; o < 32; ++o) + pos += s_sb[rr][o] < sb; + if (sel) + dst[pos] = bj; + int const cnt = __popc(__ballot_sync(0xFFFFFFFFu, sel)); + if (lane >= cnt && lane < kTopK) + dst[lane] = -1; + } + else if (c <= kCandCap) + { + for (int j = lane; j < c; j += 32) + { + uint32_t const kj = cand_key[rr * kCandCap + j]; + int const bj = cand_blk[rr * kCandCap + j]; + int rank = 0; + for (int o = 0; o < c; ++o) + { + uint32_t const ko = cand_key[rr * kCandCap + o]; + rank += (ko > kj) || (ko == kj && cand_blk[rr * kCandCap + o] < bj); + } + if (rank < kTopK) + { + s_sel_blk[rr][rank] = bj; + s_sel_neg[rr][rank] = kj == neg_key; + } + } + __syncwarp(); + bool const has = lane < kTopK && !s_sel_neg[rr][lane & 15]; + int const sk = has ? s_sel_blk[rr][lane & 15] : 0x7FFFFFFF; + int pos = 0; +#pragma unroll + for (int l = 0; l < kTopK; ++l) + { + int const o = __shfl_sync(0xFFFFFFFFu, sk, l); + pos += (o < sk) || (o == sk && l < lane); + } + if (lane < kTopK) + dst[pos] = sk == 0x7FFFFFFF ? -1 : sk; + } + else + { + // Massive exact ties: exact radix select over the staged keys. + int const nrr = max(0, min(__ldcg(p.n_valid + t), p.max_tiles)); + uint32_t const* kr = reinterpret_cast(stages) + rr * ldk; + uint32_t* hist = reinterpret_cast(stages) + rows * ldk + rr * 256; + radix_emit(kr, nrr, hist, dst, neg_key); + } + } +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +// TMA descriptor of the [pages, 128, 128] e4m3 index-K pool (one page per box, 128-byte swizzle), cached per +// (base, pages, strides). +bool make_page_map(CUtensorMap& map, void const* base, int64_t pages, int64_t stride_row, int64_t stride_page) +{ + using Key = std::tuple; + static std::mutex mu; + static std::map cache; + Key const key{reinterpret_cast(base), pages, stride_row, stride_page}; + std::lock_guard lock(mu); + auto it = cache.find(key); + if (it != cache.end()) + { + map = it->second; + return true; + } + cuuint64_t dims[3] + = {static_cast(kHeadDim), static_cast(kBlockK), static_cast(pages)}; + cuuint64_t strides[2] = {static_cast(stride_row), static_cast(stride_page)}; + cuuint32_t box[3] = {kHeadDim, kBlockK, 1}; + cuuint32_t estr[3] = {1, 1, 1}; + CUresult const res = tensorrt_llm::common::CUDADriverWrapper::getInstance()->cuTensorMapEncodeTiled(&map, + CU_TENSOR_MAP_DATA_TYPE_UINT8, 3, const_cast(base), dims, strides, box, estr, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_L2_256B, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + if (res != CUDA_SUCCESS) + { + return false; + } + if (cache.size() > 64) + { + cache.clear(); + } + cache.emplace(key, map); + return true; +} + +constexpr size_t kSmemBytes = 1024 + kNStage * kStageBytes + kEpiBytes + 2 * kNStage * 8; + +cudaError_t configure() +{ + return cudaFuncSetAttribute(m3_fused_score_select_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(kSmemBytes)); +} + +} // namespace + +cudaError_t invokeMinimaxM3Tp8FusedScoreSelect(MinimaxM3Tp8FusedScoreSelectParams const& params, void const* kPages, + int64_t numPages, int64_t kRowStride, int64_t kPageStride, int32_t batch, int32_t splits, bool usePdl, + cudaStream_t stream) +{ + if (batch < 1 || batch > kMinimaxM3Tp8FusedScoreSelectMaxBatch || splits < 1 + || (params.max_tiles + splits - 1) / splits > kMaxLocal) + { + return cudaErrorInvalidValue; + } + CUtensorMap map; + if (!make_page_map(map, kPages, numPages, kRowStride, kPageStride)) + { + return cudaErrorInvalidValue; + } + // Kernel attributes are set once per process, on the first launch (outside CUDA-graph capture: warmup runs the + // kernel eagerly first). + static cudaError_t const configured = configure(); + if (configured != cudaSuccess) + { + return configured; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(splits, batch); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = kSmemBytes; + 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_fused_score_select_kernel, map, params); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Tp8/fusedScoreSelect.h b/cpp/tensorrt_llm/kernels/minimaxM3Tp8/fusedScoreSelect.h new file mode 100644 index 000000000000..ad92b71bac21 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Tp8/fusedScoreSelect.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 + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +//! Largest number of requests (decode spans) one fused score + select launch takes. +constexpr int32_t kMinimaxM3Tp8FusedScoreSelectMaxBatch = 2; +//! Largest decode query length (query rows per request). +constexpr int32_t kMinimaxM3Tp8FusedScoreSelectMaxDql = 4; +//! Largest number of selected rows per request (num_kv_heads * dql). +constexpr int32_t kMinimaxM3Tp8FusedScoreSelectMaxRows = 8; +//! Largest number of blocks one CTA scores (ceil(max_tiles / splits) must not exceed it). +constexpr int32_t kMinimaxM3Tp8FusedScoreSelectMaxLocalBlocks = 128; +//! Bytes of the CTA's TMA stage ring, reused by the select: rows * ceil4(max_tiles) * 4 + rows * 1024 must fit. +constexpr int64_t kMinimaxM3Tp8FusedScoreSelectRingBytes = 12 * 128 * 128; + +//! Operands of the MiniMax-M3 TP8 fused index score + top-16 block select. Strides are in elements. +struct MinimaxM3Tp8FusedScoreSelectParams +{ + uint8_t const* q; // e4m3 [total_q, heads, 128] + int64_t q_st, q_sh; + int32_t const* bt; // [batch, bt_width] + int64_t bt_stride; + int bt_width; + int32_t const* seq_lens; // [batch], including this step's dql query rows + int32_t const* n_valid; // [total_q] causal valid-block counts + float* scr; // [batch][rows][>= ceil4(max_tiles)] block scores + int64_t scr_b, scr_r; + int32_t* out; // logical [total_q, kv_heads, 16] + int64_t out_t, out_h; + int32_t* counters; // [batch], zero at rest + int heads, kv_heads, dql, max_tiles, init_blocks, local_blocks; +}; + +//! One launch of index scoring and top-16 selection for at most 2 decode requests. Grid (splits, batch), one CTA of +//! 8 consumer warps + 1 TMA producer warp per SM: CTA (s, b) TMA-loads the e4m3 [128, 128] index-K pages of blocks +//! s, s + splits, ... of request b (12-stage ring; the pages holding only tokens older than this step are requested +//! before the grid dependency wait), scores them against the request's q rows on FP16 tensor cores, and stores each +//! block's score (max over its causal positions, amax over each KV head's index heads) to the scratch. The last CTA +//! of a request (per-request arrival counter) releases PDL dependents, resets that counter to zero, and selects the +//! top 16 of every (query row, KV head) row exactly like minimax_m3_select_blocks (init / local forcing, ties to the +//! smaller block, -inf picks dropped, ascending ids with a -1 tail). +//! +//! kPages is the e4m3 index-K page pool [numPages, 128, 128] with the given byte strides (row stride and page stride +//! multiples of 16, base 16-byte aligned); its TMA descriptor is encoded on first use and cached per (base, pages, +//! strides). The device code is SM100-family only; on other architectures the kernel compiles to an empty body. +cudaError_t invokeMinimaxM3Tp8FusedScoreSelect(MinimaxM3Tp8FusedScoreSelectParams const& params, void const* kPages, + int64_t numPages, int64_t kRowStride, int64_t kPageStride, int32_t batch, int32_t splits, bool usePdl, + cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Tp8/sparseDecodeLc8.cu b/cpp/tensorrt_llm/kernels/minimaxM3Tp8/sparseDecodeLc8.cu new file mode 100644 index 000000000000..5d5b0bf605d2 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Tp8/sparseDecodeLc8.cu @@ -0,0 +1,610 @@ +/* + * 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. + +#include "tensorrt_llm/kernels/minimaxM3Tp8/sparseDecodeLc8.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +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]; + float part[kHalfDim]; + float red[4][kHeads8]; + float rowm[kHeads8]; + float rowl[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) + +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"); + 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 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); + } + } + load_q(s, p, tok, tid); + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); + } + int blk = static_cast(rank) < real_topk ? p.topk[tok * p.t_tok + rank * p.t_k] : -1; + bool const active = blk >= 0; + if constexpr (PREWAIT) + { + __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(); + if (tid < kHeads8) + 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[2 * qd]; + m_row[1] = s.rowm[2 * qd + 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_row[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(); // rowm reads 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(); + if (tid < kHeads8) + 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[2 * qd]; + l_row[1] = s.rowl[2 * qd + 1]; + } + + // 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 (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] + 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]); + } + } + 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; + float lse_c[kTopk]; + float lse_max = kNegInf; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + { + lse_c[c] = s.recv_lse[c]; + lse_max = fmaxf(lse_max, lse_c[c]); + } + lse_max = lse_max == kNegInf ? 0.F : lse_max; + float denom = 0.F; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + { + lse_c[c] = ex2(lse_c[c] - lse_max); + denom += lse_c[c]; + } + denom = denom > 0.F ? denom : 1.F; + 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] * __fdiv_rn(lse_c[cc], denom); + } + if (hc) + s.part[dl] = o; + __syncthreads(); + if (!hc) + { + o += s.part[dl]; + 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) +{ + return prewait ? launch(prm, num_tokens, use_pdl, stream) + : launch(prm, num_tokens, use_pdl, stream); +} + +} // namespace + +cudaError_t invokeMinimaxM3Tp8SparseDecodeLc8(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 kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Tp8/sparseDecodeLc8.h b/cpp/tensorrt_llm/kernels/minimaxM3Tp8/sparseDecodeLc8.h new file mode 100644 index 000000000000..ad65a577eab8 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Tp8/sparseDecodeLc8.h @@ -0,0 +1,75 @@ +/* + * 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 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; +}; + +//! 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 invokeMinimaxM3Tp8SparseDecodeLc8(MinimaxM3Tp8SparseDecodeLc8Params 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/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 3345690fa796..84c8ea6d9407 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -84,6 +84,8 @@ add_th_common_family( minimaxM3SparseDecodeLcOp.cpp minimaxM3SpecTailOp.cpp minimaxM3TargetHeadOp.cpp + minimaxM3Tp8FusedScoreSelectOp.cpp + minimaxM3Tp8SparseDecodeLc8Op.cpp stepInputsScatterOp.cpp mlaRopeInplaceOp.cpp relativeAttentionBiasOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3Tp8FusedScoreSelectOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3Tp8FusedScoreSelectOp.cpp new file mode 100644 index 000000000000..1d28f41fdc7a --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3Tp8FusedScoreSelectOp.cpp @@ -0,0 +1,173 @@ +/* + * 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/minimaxM3Tp8/fusedScoreSelect.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kHeadDim = 128; +constexpr int64_t kBlockK = 128; +constexpr int64_t kTopK = 16; + +} // namespace + +//! MiniMax-M3 TP8 fused index score + top-16 block select for decode spans of at most 2 requests. Equals +//! minimax_m3_select_blocks applied to the block scores of cute_dsl_minimax_m3_index_decode_score. Grid (splits, +//! batch) with splits = max(SM count / batch, ceil(max_tiles / 128), 1). +//! +//! @param idxQ E4M3 [total_q, heads, 128] (total_q = batch * dql, heads * 4 <= 8), unit last-dim stride, row and head +//! strides multiples of 4 bytes, 4-byte aligned. +//! @param idxKPages E4M3 [num_pages, 128, 128], unit last-dim stride, page and row strides multiples of 16 bytes, +//! 16-byte aligned. +//! @param blockTable Int32 [batch, >= blocks], unit last-dim stride. +//! @param seqLens Int32 [batch] (1 or 2 requests), including this step's dql query rows. +//! @param nValid Contiguous int32 [total_q] per-token causal valid-block counts. +//! @param out Int32, logical [total_q, num_kv_heads, 16] with unit last-dim stride (a permuted head-major backing is +//! fine). Every entry is written: ascending block ids with a -1 tail. +//! @param scratch Contiguous float32 [>= batch, >= num_kv_heads * dql, L] with L >= max_tiles and L % 4 == 0, 16-byte +//! aligned; [batch, 8, ceil4(max_tiles)] covers every supported geometry. Contents need no initialization. +//! @param counters Contiguous int32 with >= batch elements, all zero before the first call. The kernel returns each +//! request's counter to zero before it exits, so the buffer stays reusable across calls and graph replays without a +//! reset (as long as launches sharing it are stream-ordered). +//! @param maxTiles Worst-case block count n_valid is clamped to; num_kv_heads * dql * ceil4(max_tiles) * 4 + +//! num_kv_heads * dql * 1024 must not exceed 196608 bytes. +void minimaxM3Tp8FusedScoreSelect(torch::Tensor const& idxQ, torch::Tensor const& idxKPages, + torch::Tensor const& blockTable, torch::Tensor const& seqLens, torch::Tensor const& nValid, torch::Tensor& out, + torch::Tensor& scratch, torch::Tensor& counters, int64_t decodeQueryLen, int64_t numKvHeads, int64_t initBlocks, + int64_t localBlocks, int64_t maxTiles, bool usePdl) +{ + namespace tk = tensorrt_llm::kernels; + TORCH_CHECK(idxQ.is_cuda(), "minimax_m3_tp8_fused_score_select: idx_q must be a CUDA tensor"); + c10::cuda::CUDAGuard const deviceGuard{idxQ.device()}; + auto const* props = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(props->major >= 10, "minimax_m3_tp8_fused_score_select requires an SM100-family GPU"); + + TORCH_CHECK(idxQ.dim() == 3 && idxQ.size(2) == kHeadDim && idxQ.stride(2) == 1 && idxQ.element_size() == 1, + "minimax_m3_tp8_fused_score_select: idx_q must be one-byte [total_q, heads, 128] with unit last-dim stride"); + TORCH_CHECK( + idxQ.stride(0) % 4 == 0 && idxQ.stride(1) % 4 == 0 && (reinterpret_cast(idxQ.data_ptr()) & 3) == 0, + "minimax_m3_tp8_fused_score_select: idx_q strides and base must be 4-byte aligned"); + TORCH_CHECK(idxKPages.is_cuda() && idxKPages.element_size() == 1 && idxKPages.dim() == 3 + && idxKPages.size(1) == kBlockK && idxKPages.size(2) == kHeadDim && idxKPages.stride(2) == 1, + "minimax_m3_tp8_fused_score_select: idx_k_pages must be one-byte [num_pages, 128, 128] with unit last-dim " + "stride"); + TORCH_CHECK(idxKPages.stride(1) % 16 == 0 && idxKPages.stride(0) % 16 == 0 + && (reinterpret_cast(idxKPages.data_ptr()) & 15) == 0, + "minimax_m3_tp8_fused_score_select: idx_k_pages strides and base must be 16-byte aligned"); + TORCH_CHECK(blockTable.scalar_type() == torch::kInt32 && blockTable.dim() == 2 && blockTable.stride(1) == 1, + "minimax_m3_tp8_fused_score_select: block_table must be int32 [batch, blocks] with unit last-dim stride"); + TORCH_CHECK( + seqLens.scalar_type() == torch::kInt32 && nValid.scalar_type() == torch::kInt32 && nValid.is_contiguous(), + "minimax_m3_tp8_fused_score_select: seq_lens and n_valid must be int32, n_valid contiguous"); + TORCH_CHECK(out.scalar_type() == torch::kInt32 && out.dim() == 3 && out.size(2) == kTopK && out.stride(2) == 1, + "minimax_m3_tp8_fused_score_select: out must be int32 [total_q, num_kv_heads, 16] with unit last-dim stride"); + + int64_t const batch = seqLens.size(0); + int64_t const totalQ = idxQ.size(0); + int64_t const heads = idxQ.size(1); + int64_t const rows = decodeQueryLen * numKvHeads; + TORCH_CHECK(batch >= 1 && batch <= tk::kMinimaxM3Tp8FusedScoreSelectMaxBatch, + "minimax_m3_tp8_fused_score_select takes 1 to ", tk::kMinimaxM3Tp8FusedScoreSelectMaxBatch, " requests, got ", + batch); + TORCH_CHECK(decodeQueryLen >= 1 && decodeQueryLen <= tk::kMinimaxM3Tp8FusedScoreSelectMaxDql + && totalQ == batch * decodeQueryLen, + "minimax_m3_tp8_fused_score_select: decode_query_len must be in [1, 4] with total_q = batch * " + "decode_query_len"); + TORCH_CHECK(numKvHeads >= 1 && heads % numKvHeads == 0 && heads * tk::kMinimaxM3Tp8FusedScoreSelectMaxDql <= 8 + && rows <= tk::kMinimaxM3Tp8FusedScoreSelectMaxRows, + "minimax_m3_tp8_fused_score_select: unsupported head geometry"); + TORCH_CHECK(out.size(0) == totalQ && out.size(1) == numKvHeads, + "minimax_m3_tp8_fused_score_select: out must be [total_q, num_kv_heads, 16]"); + TORCH_CHECK(nValid.numel() >= totalQ, "minimax_m3_tp8_fused_score_select: n_valid needs total_q elements"); + TORCH_CHECK(maxTiles >= 1, "minimax_m3_tp8_fused_score_select: max_tiles must be positive"); + + TORCH_CHECK(scratch.is_cuda() && scratch.scalar_type() == torch::kFloat32 && scratch.dim() == 3 + && scratch.is_contiguous() && scratch.size(0) >= batch && scratch.size(1) >= rows + && scratch.size(2) >= maxTiles && scratch.stride(1) % 4 == 0 + && (reinterpret_cast(scratch.data_ptr()) & 15) == 0, + "minimax_m3_tp8_fused_score_select: scratch must be contiguous 16-byte aligned float32 [>= batch, >= rows, " + "L] with L >= max_tiles and L % 4 == 0"); + TORCH_CHECK(counters.is_cuda() && counters.scalar_type() == torch::kInt32 && counters.is_contiguous() + && counters.numel() >= batch, + "minimax_m3_tp8_fused_score_select: counters must be contiguous int32 with >= batch elements"); + TORCH_CHECK(rows * ((maxTiles + 3) & ~int64_t{3}) * 4 + rows * 1024 <= tk::kMinimaxM3Tp8FusedScoreSelectRingBytes, + "minimax_m3_tp8_fused_score_select: select scratch does not fit the stage ring"); + + int64_t const splits = std::max({props->multiProcessorCount / batch, + (maxTiles + tk::kMinimaxM3Tp8FusedScoreSelectMaxLocalBlocks - 1) + / tk::kMinimaxM3Tp8FusedScoreSelectMaxLocalBlocks, + int64_t{1}}); + TORCH_CHECK((maxTiles + splits - 1) / splits <= tk::kMinimaxM3Tp8FusedScoreSelectMaxLocalBlocks, + "minimax_m3_tp8_fused_score_select: too few splits"); + + tk::MinimaxM3Tp8FusedScoreSelectParams p{}; + p.q = reinterpret_cast(idxQ.data_ptr()); + p.q_st = idxQ.stride(0); + p.q_sh = idxQ.stride(1); + p.bt = blockTable.data_ptr(); + p.bt_stride = blockTable.stride(0); + p.bt_width = static_cast(blockTable.size(1)); + p.seq_lens = seqLens.data_ptr(); + p.n_valid = nValid.data_ptr(); + p.scr = scratch.data_ptr(); + p.scr_b = scratch.stride(0); + p.scr_r = scratch.stride(1); + p.out = out.data_ptr(); + p.out_t = out.stride(0); + p.out_h = out.stride(1); + p.counters = counters.data_ptr(); + p.heads = static_cast(heads); + p.kv_heads = static_cast(numKvHeads); + p.dql = static_cast(decodeQueryLen); + p.max_tiles = static_cast(maxTiles); + p.init_blocks = static_cast(initBlocks); + p.local_blocks = static_cast(localBlocks); + + auto const stream = at::cuda::getCurrentCUDAStream(idxQ.get_device()); + cudaError_t const err + = tk::invokeMinimaxM3Tp8FusedScoreSelect(p, idxKPages.data_ptr(), idxKPages.size(0), idxKPages.stride(1), + idxKPages.stride(0), static_cast(batch), static_cast(splits), usePdl, stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_tp8_fused_score_select launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_tp8_fused_score_select(Tensor idx_q, Tensor idx_k_pages, Tensor block_table, Tensor seq_lens, " + "Tensor n_valid, Tensor(a!) out, Tensor(b!) scratch, Tensor(c!) counters, int decode_query_len, " + "int num_kv_heads, int init_blocks, int local_blocks, int max_tiles, bool use_pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_tp8_fused_score_select", &tensorrt_llm::torch_ext::minimaxM3Tp8FusedScoreSelect); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3Tp8SparseDecodeLc8Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3Tp8SparseDecodeLc8Op.cpp new file mode 100644 index 000000000000..3f40efdb9bb6 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3Tp8SparseDecodeLc8Op.cpp @@ -0,0 +1,139 @@ +/* + * 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/minimaxM3Tp8/sparseDecodeLc8.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_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 minimaxM3Tp8SparseDecodeLc8(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) +{ + TORCH_CHECK(q.is_cuda(), "minimax_m3_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_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_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_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_tp8_sparse_decode_lc8: topk must be int32 [T, 16]"); + TORCH_CHECK(blockTable.scalar_type() == torch::kInt32 && seqLens.scalar_type() == torch::kInt32, + "minimax_m3_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_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_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::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); + + auto const stream = at::cuda::getCurrentCUDAStream(q.get_device()); + cudaError_t const err = tensorrt_llm::kernels::invokeMinimaxM3Tp8SparseDecodeLc8( + prm, static_cast(q.size(0)), linearSoftmax, mxfp8, prewait, usePdl, stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_tp8_sparse_decode_lc8 launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_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) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_tp8_sparse_decode_lc8", &tensorrt_llm::torch_ext::minimaxM3Tp8SparseDecodeLc8); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/__init__.py new file mode 100644 index 000000000000..2e340ff7284d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/__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 NVFP4 / sm_103 / tp8.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/dispatch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/dispatch.py new file mode 100644 index 000000000000..26df46630059 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/dispatch.py @@ -0,0 +1,527 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-step fast paths of this target: the rows each one takes, the catalog entry that certifies them, and +what happens when one cannot run. + +One target serves every concurrency with the same code; which fast path a step takes is a pure function of its +shape, written here once (``plan_for``, ``moe_path``, ``router_path``, ``attention_path``, ``linear_path``) and +called by the forward. Every row range is a named constant below, and ``FAST_PATHS`` pairs each with the catalog +entry whose ``CELLS`` certify it: ``test_modeling_v2_m3_dispatch.py`` checks that every row count (or shape key) a +range covers is one of the entry's cells. + +Two kinds of not taking a fast path, kept apart: + +* **uncovered**: the step is outside every covered range -- prefill, mixed steps, generation wider than the ranges, + an eager decode step (the GEMV dispatch is a decode-graph one), a long-context request, an Eagle3 capture layer + that runs the generic layer by design. ``uncovered`` logs the reason once per (path, reason); the caller runs the + generic path in every mode. +* **missed**: the step is inside a covered range but a precondition is missing -- a workspace, a push descriptor or + a compiled engine not built, or a kernel declining operands it is certified for. Under ``modeling_v2='require'`` + ``missed`` raises ``FastPathUnavailable`` naming the path, the rows and the precondition; otherwise it logs once + per (path, reason) and the caller falls back. + +Compiled kernels (the MoE engines, the decode GEMVs with their push variants, the index scorer) are built by +``build_targets`` 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, kernel modules) 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``) rather +than a silent capture of the fallback. +""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Callable, Dict, FrozenSet, Hashable, Iterable, Mapping, Optional, Tuple + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper +from tensorrt_llm.logger import logger + +from .ops.attention.m3_dense_decode import m3_dense_decode +from .ops.attention.m3_fused_score_select import ( + MAX_REQUESTS as FUSED_SCORE_SELECT_KERNEL_MAX_REQUESTS, +) +from .ops.attention.m3_fused_score_select import m3_fused_score_select +from .ops.attention.m3_index_score import m3_index_score +from .ops.attention.m3_producer import m3_producer +from .ops.attention.m3_select_long import MAX_BLOCKS as SELECT_LONG_MAX_BLOCKS +from .ops.attention.m3_select_long import m3_select_long +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_mc_allreduce_norm import m3_mc_allreduce_norm +from .ops.comm.m3_mc_finalize_allreduce_norm import m3_mc_finalize_allreduce_norm +from .ops.gemm.m3_draft_lowm import TACTICS as DRAFT_LOWM_TACTICS +from .ops.gemm.m3_draft_lowm import m3_draft_lowm +from .ops.gemm.m3_mx_gemv import DENSE_DOWN, DENSE_GATE_UP, O_PROJ, QKV, QKV_SPARSE, m3_mx_gemv +from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 +from .ops.moe.m3_moe_engine import m3_moe_engine +from .ops.moe.m3_router_gemv import m3_router_gemv +from .ops.moe.m3_router_gemv_fp4 import m3_router_gemv_fp4 +from .ops.moe.m3_router_tc import m3_router_tc +from .ops.norm.m3_embedding_firstnorm import m3_embedding_firstnorm +from .ops.spec.m3_draft_input import m3_draft_input +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_shard_accept import m3_shard_accept + +# -------------------------------------------------------------------------------------------------------------------- +# Row ranges +# -------------------------------------------------------------------------------------------------------------------- + +# The largest generation step the decode target takes: 8 requests verifying 1 + 3 drafts and their padding rows. +DECODE_MAX_TOKENS = 48 + + +@dataclass(frozen=True) +class DecodePlan: + """The layer-boundary collectives of one decode token bucket (GB300 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 (kernels/m3_mc_allreduce.BUCKETS). +DECODE_PLANS = ( + DecodePlan("tiny", 8, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), + DecodePlan("small", 16, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), + DecodePlan("verify", DECODE_MAX_TOKENS, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), +) + +# 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. Two requests are neutral (the grid is unchanged) and keep the step's bound. +DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({1: 16384}) +# The same cap on the Eagle3 drafter layer's trtllm-gen decode (8 query heads over one KV head per rank): GB300, one +# request, 63.7K cached tokens: 22.5 -> 11.7 us at one query row, 24.3 -> 12.5 us at four; 188K tokens: 28.2 -> 20.2 us. +# Two requests are flat and keep the step's bound. +DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({1: 16384}) + +# The sparse layers' NVFP4 QKV + indexer producer (ops/attention/m3_producer, every step its geometry fits) triggers its +# dependent kernel on entry, so the kernel after it launches and stages its pre-wait loads while the producer runs. +PRODUCER_EARLY_TRIGGER = True + +# The fused embedding + first RMSNorm kernel: one owner per row count and index dtype. +EMBED_FIRSTNORM_MAX_TOKENS = 8 +# 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 = 16 +AR_PUSH_MODE = 2 +# The 1..8-token MoE engine (router, routed experts at the TP8 width 384 and the shared expert in one launch) pushes +# its sum into the MoE boundary's all-reduce (prepushed mode 2). +MOE_ENGINE_MAX_TOKENS = 8 +MOE_PUSH_MAX_TOKENS = MOE_ENGINE_MAX_TOKENS +# 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. +MC_FINALIZE_MIN_TOKENS = MOE_ENGINE_MAX_TOKENS + 1 +# The router: the fp32 GEMV at one row, the tensor-core router (fused with the routed experts' NVFP4 activation +# quantization) at 2..48 rows. +ROUTER_GEMV_MAX_TOKENS = 1 +ROUTER_TC_MAX_TOKENS = 48 +# 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 (kernels/m3_shard_accept), at the configured draft length. The lm_head +# is the in-tree BF16 GEMM: at the TP8 shard (25008 rows) a coded copy of the head saves no time. +SHARD_ACCEPT_DRAFT_LEN = 3 +SHARD_ACCEPT_MAX_REQUESTS = DECODE_MAX_TOKENS // (SHARD_ACCEPT_DRAFT_LEN + 1) +# 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 (spec_worker). At TP8 it forks beside the verify head instead: forked early its +# GEMM lands beside the target tail it was meant to hide under. +FC_EARLY_MAX_ROWS = 0 +# 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 (kernels/m3_spec_tail.draft_input), the drafter layer's in-tree attention, MLP and +# all-reduces, and one kernel for the BF16 draft head, the cross-rank argmax and the pass's bookkeeping +# (kernels/m3_draft_step), at the configured draft length. The head kernel takes one row per request, at most 8, and +# streams the head's BF16 weight fragments. +DRAFT_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN +DRAFT_STEP_MAX_REQUESTS = DRAFT_STEP_KERNEL_MAX_TOKENS +# The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows +# (kernels/triton_sparse_decode.MQ_MIN_TOKENS). +SPARSE_DECODE_MQ_MIN_TOKENS = 17 +# Decode spans of at most this many requests score and select their top-16 blocks in one launch +# (ops/attention/m3_fused_score_select), which loads the older index-K pages before its grid wait; larger spans run +# the index scorer and the long-row selector. +FUSED_SCORE_SELECT_MAX_REQUESTS = FUSED_SCORE_SELECT_KERNEL_MAX_REQUESTS +# 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 +LONG_SELECT_MAX_BLOCKS = SELECT_LONG_MAX_BLOCKS +LONG_SELECT_KEY_WIDTHS = (2, 4, 9) +# The long-row select releases the sparse decode after it right after its own grid wait up to this many query rows; +# wider steps keep the trigger at the select's exit (the early-resident attention CTAs crowd the select grid). +SELECT_EARLY_MAX_ROWS = 8 +# The 8-head long-context sparse decode (ops/attention/m3_sparse_decode_lc8) stages the request length, its query rows +# and the request's block-table row before its grid wait up to this many query rows; above, the staged row's shared +# memory costs occupancy and the kernel stages only after the wait (same results bit for bit). +LC8_PREWAIT_MAX_ROWS = 8 +# The MSA index scorer (ops/attention/m3_index_score) prefetches the older index-K pages into L2 before its grid wait +# at every batch it runs (more than FUSED_SCORE_SELECT_MAX_REQUESTS requests); it never runs its one-request +# early-issue schedule here. +INDEX_SCORE_EARLY_MAX_REQUESTS = 0 + +# 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). +_ALL_ROWS = frozenset(range(1, DECODE_MAX_TOKENS + 1)) +MX_GEMV_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { + QKV_SPARSE: _ALL_ROWS, + QKV: _ALL_ROWS, + O_PROJ: _ALL_ROWS, + DENSE_GATE_UP: _ALL_ROWS, + DENSE_DOWN: _ALL_ROWS, +} +# 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: 48, + DENSE_GATE_UP: 16, + DENSE_DOWN: 32, +} +# The Eagle3 drafter's BF16 linears (qkv, o_proj, gate_up, down, fc) through FlashInfer's fixed-tactic low-M kernels +# (ops/gemm/m3_draft_lowm): per (N, K), the rows its tactic table covers; more rows keep the in-tree GEMM. +DRAFT_LOWM_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { + nk: frozenset(range(1, table[-1][0] + 1)) for nk, table in DRAFT_LOWM_TACTICS.items() +} + +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, +) +# The long-context sparse decode: "prewait" up to LC8_PREWAIT_MAX_ROWS query rows, "plain" above. +SPARSE_DECODE_LC8 = RowRange( + "m3_sparse_decode_lc8", + frozenset({"prewait", "plain"}), + 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. +PRODUCER = RowRange("m3_producer", frozenset({"nvfp4_tp8"}), m3_producer, _const("nvfp4_tp8")) +# Requests of a pure decode span whose blocks are scored and selected in one launch. +FUSED_SCORE_SELECT = RowRange( + "m3_fused_score_select", + _rows(1, FUSED_SCORE_SELECT_MAX_REQUESTS), + m3_fused_score_select, + _spec_rows("batch"), +) +INDEX_SCORE = RowRange( + "m3_index_score", + frozenset({"prefetch"}), + m3_index_score, + lambda spec: () if spec["early"] else ("prefetch",), +) +LONG_SELECT = RowRange( + "long_select", frozenset(LONG_SELECT_KEY_WIDTHS), m3_select_long, _long_select_widths +) +MC_ALLREDUCE = RowRange( + "mc_allreduce_norm", + _rows(1, DECODE_MAX_TOKENS), + 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), +) +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, DECODE_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") +) +ROUTER_GEMV = RowRange( + "m3_router_gemv", _rows(1, ROUTER_GEMV_MAX_TOKENS), m3_router_gemv, _spec_rows("m") +) +ROUTER_GEMV_FP4 = RowRange( + "m3_router_gemv_fp4", _rows(1, ROUTER_GEMV_MAX_TOKENS), m3_router_gemv_fp4, _spec_rows("m") +) +ROUTER_TC = RowRange( + "m3_router_tc", + _rows(ROUTER_GEMV_MAX_TOKENS + 1, ROUTER_TC_MAX_TOKENS), + m3_router_tc, + _spec_rows("m"), +) +# 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_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 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 (), +) + +# (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, + DENSE_DECODE_SPLIT, + SPARSE_DECODE_PRODUCER, + SPARSE_DECODE_LC8, + PRODUCER, + FUSED_SCORE_SELECT, + INDEX_SCORE, + LONG_SELECT, + MC_ALLREDUCE, + ATTN_PUSH, + MOE_PUSH, + MC_FINALIZE, + MOE_ENGINE, + ROUTER_GEMV, + ROUTER_GEMV_FP4, + ROUTER_TC, + MX_GEMV, + MX_GEMV_BF16, + DRAFT_LOWM, + SHARD_ACCEPT, + DRAFT_STEP, + 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 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 + + +def router_path(num_tokens: int) -> Optional[str]: + """The router kernel when the MoE runs unfused: the fp32 GEMV, the tensor-core router, or None (built-in gate).""" + if ROUTER_GEMV.covers(num_tokens): + return ROUTER_GEMV.path + if ROUTER_TC.covers(num_tokens): + return ROUTER_TC.path + return None + + +MX_GEMV_PATH = MX_GEMV.path +MX_GEMV_BF16_PATH = MX_GEMV_BF16.path +# 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) -> str: + """The decode-graph MXFP8 GEMM of an [N, K] projection on ``num_tokens`` rows. + + ``prequant``: the producer published the activation in MXFP8. + """ + 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 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() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/__init__.py new file mode 100644 index 000000000000..093ae922e88a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kernels of the MiniMax-M3 NVFP4 / sm_103 / tp8 target.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/long_select.py new file mode 100644 index 000000000000..4fb5064fa981 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/long_select.py @@ -0,0 +1,62 @@ +# 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_select_long``, 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 + + +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, +) -> 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 (``dispatch.SELECT_EARLY_MAX_ROWS``). + """ + 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( + max_score_kv, + n_valid_blocks.contiguous(), + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_attn_lc8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_attn_lc8.py new file mode 100644 index 000000000000..98b601471d10 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_attn_lc8.py @@ -0,0 +1,110 @@ +# 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 ``dispatch.LC8_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 (``dispatch.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 .. import dispatch +from ..ops.attention.m3_sparse_decode_lc8 import NUM_HEADS, WIDTH, m3_sparse_decode_lc8 + + +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, +): + """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. + """ + 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 <= dispatch.LC8_PREWAIT_MAX_ROWS, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_draft_lowm.py new file mode 100644 index 000000000000..356c26639947 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_draft_lowm.py @@ -0,0 +1,121 @@ +# 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 (``dispatch.DRAFT_LOWM``). + +``install`` swaps the linear method of every bias-free BF16 ``Linear`` of the drafter (qkv, o_proj, gate_up, down, the +fc; not the head) for ``DraftLowMLinearMethod``: a covered row count runs ``ops/gemm/m3_draft_lowm`` at the tactic +of its (N, K) table, every other call the in-tree method. ``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 Optional, Set, Tuple + +import torch +from torch import nn + +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import Linear, UnquantizedLinearMethod + +from .. import dispatch +from ..ops.gemm.m3_draft_lowm import m3_draft_lowm, tactic_for + +# (N, K, M, PDL) of every call prebuild compiled. +_BUILT: Set[Tuple[int, int, int, bool]] = set() + + +class DraftLowMLinearMethod(UnquantizedLinearMethod): + """The in-tree BF16 linear method with the drafter's covered rows on the fixed-tactic low-M kernels.""" + + def apply(self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor]): + out = draft_linear(module.weight, input, bias) + if out is not None: + return out + return super().apply(module, input, bias) + + +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 ``dispatch.DRAFT_LOWM`` covers it; None for the in-tree + GEMM.""" + if bias is not None or x.dim() != 2: + return None + n, k = weight.shape + m = int(x.shape[0]) + if not dispatch.DRAFT_LOWM.covers(((n, k), m)): + dispatch.uncovered( + dispatch.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. + dispatch.uncovered( + dispatch.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: + dispatch.missed(dispatch.DRAFT_LOWM.path, m, "the activation or weight is not BF16") + return None + if (n, k, m, use_pdl) not in _BUILT: + dispatch.missed( + dispatch.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) + out = torch.empty((m, n), dtype=x.dtype, device=x.device) + return m3_draft_lowm(x, weight, out, kind, tactic, use_pdl) + + +def _drafter_linears(draft_model: nn.Module): + for module in draft_model.modules(): + if ( + isinstance(module, Linear) + and not isinstance(module, LMHead) + and type(module.quant_method) is UnquantizedLinearMethod + and getattr(module, "bias", None) is None + and module.weight.dtype == torch.bfloat16 + ): + yield module + + +def install(draft_model: nn.Module) -> int: + """Put every covered drafter linear on ``DraftLowMLinearMethod``; returns how many were installed.""" + count = 0 + for module in _drafter_linears(draft_model): + n, k = module.weight.shape + if any( + dispatch.DRAFT_LOWM.covers(((n, k), m)) + for m in range(1, dispatch.DECODE_MAX_TOKENS + 1) + ): + module.quant_method = DraftLowMLinearMethod() + count += 1 + return count + + +def prebuild(draft_model: nn.Module) -> int: + """Compile (and run once) every covered row count of every installed drafter linear, outside capture, on the + current stream. Returns the number of calls compiled.""" + use_pdl = get_env_enable_pdl() + built = 0 + shapes = { + tuple(module.weight.shape): module.weight + for module in draft_model.modules() + if isinstance(module, Linear) and isinstance(module.quant_method, DraftLowMLinearMethod) + } + for (n, k), weight in sorted(shapes.items()): + for nk, m in sorted(key for key in dispatch.DRAFT_LOWM.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) + _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_103__tp8/kernels/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_draft_step.py new file mode 100644 index 000000000000..42aee2cb9041 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_draft_step.py @@ -0,0 +1,150 @@ +# 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 + ``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. +* 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 .. import dispatch +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 .m3_shard_accept import AcceptWorkspace, create_workspace + + +@dataclass(eq=False) +class DraftHead: + """The streamed copy of one rank's draft lm_head shard and its first global id: ``wf`` is the fragment-ordered BF16 + copy (``bf16_fragments``).""" + + weight: torch.Tensor + wf: torch.Tensor + vocab_offset: int + + +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 fragment-ordered BF16 copy of ``lm_head``'s local shard (global ids from tp_rank * rows). + 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(): + wf = bf16_fragments(weight.detach()) + return DraftHead(weight=weight, wf=wf, vocab_offset=offset) + + +_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(): + dispatch.capture_missed(dispatch.DRAFT_STEP.path, 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``), 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, + _num_sms(), + True, + None, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_fused_score_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_fused_score_select.py new file mode 100644 index 000000000000..3ef07ce6dbff --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_fused_score_select.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fused index score + top-16 select of a decode span of at most ``dispatch.FUSED_SCORE_SELECT_MAX_REQUESTS`` +requests (``ops/attention/m3_fused_score_select``), with its scratch and arrival counters. + +One scratch [2, 8, ceil4(max_tiles)] and one zeroed counter pair per (device, engine block count) serve every sparse +layer: the layers run one after another on one stream, and the kernel returns each counter to zero. They are +created on an eager decode step; a CUDA-graph capture only reads them. +""" + +from __future__ import annotations + +from typing import Dict, Optional, Tuple + +import torch + +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +from ..ops.attention.m3_fused_score_select import ( + MAX_DECODE_QUERY_LEN, + MAX_ROWS, + m3_fused_score_select, + ring_fits, + scratch_shape, +) + +_BUFFERS: Dict[Tuple[str, int], Tuple[torch.Tensor, torch.Tensor]] = {} + + +def supported( + idx_q: torch.Tensor, + idx_k_pages: torch.Tensor, + decode_query_len: int, + num_kv_heads: int, + max_tiles: int, +): + """None when the kernel takes this span's geometry, else the reason it does not.""" + if ( + idx_q.dim() != 3 + or idx_q.dtype != torch.float8_e4m3fn + or idx_k_pages.dtype != torch.float8_e4m3fn + ): + return "the index Q / index-K cache is not FP8" + heads = int(idx_q.shape[1]) + rows = decode_query_len * num_kv_heads + if ( + idx_q.shape[2] != 128 + or idx_k_pages.dim() != 3 + or tuple(idx_k_pages.shape[1:]) != (128, 128) + or not 1 <= decode_query_len <= MAX_DECODE_QUERY_LEN + or heads % num_kv_heads + or heads * MAX_DECODE_QUERY_LEN > MAX_ROWS + or rows > MAX_ROWS + ): + return f"{heads} index heads x {decode_query_len} query rows is not the kernel's geometry" + if not ring_fits(rows, max_tiles): + return f"{max_tiles} blocks of select scratch do not fit the stage ring" + return None + + +def buffers(device: torch.device, max_tiles: int) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """(scratch, counters) for this engine's block count; None during a capture that finds them missing.""" + key = (str(device), int(max_tiles)) + found = _BUFFERS.get(key) + if found is None: + if torch.cuda.is_current_stream_capturing(): + return None + found = _BUFFERS[key] = ( + torch.empty(scratch_shape(max_tiles), dtype=torch.float32, device=device), + torch.zeros(scratch_shape(max_tiles)[0], dtype=torch.int32, device=device), + ) + return found + + +def fused_score_select( + idx_q, + idx_k_pages, + block_table, + seq_lens, + n_valid, + *, + decode_query_len, + num_kv_heads, + init_blocks, + local_blocks, + max_tiles, + scratch, + counters, +): + """The [total_q, num_kv_heads, 16] table (a view of a head-major [num_kv_heads, total_q, 16] backing).""" + total_q = int(idx_q.shape[0]) + out = torch.empty((num_kv_heads, total_q, 16), dtype=torch.int32, device=idx_q.device).permute( + 1, 0, 2 + ) + return m3_fused_score_select( + idx_q, + idx_k_pages, + block_table, + seq_lens, + n_valid, + out, + scratch, + counters, + int(decode_query_len), + int(num_kv_heads), + int(init_blocks), + int(local_blocks), + int(max_tiles), + get_env_enable_pdl(), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_index_score.py new file mode 100644 index 000000000000..9748bd1fe4a5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_index_score.py @@ -0,0 +1,645 @@ +# 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/cute_dsl_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.cute_dsl_kernels.blackwell.cute_ptx_utils import ( + EVICT_FIRST, + fp8x4_to_fp16x4, + mma_sync, + simple_tma_copy, +) +from tensorrt_llm._torch.cute_dsl_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 + SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K + # BLOCK_Q must fit one warp's worth of epilogue lanes. + MAX_BLOCK_Q = 32 + 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.cute_dsl_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.cute_dsl_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_103__tp8/kernels/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_mc_allreduce.py new file mode 100644 index 000000000000..7617e9f64772 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_mc_allreduce.py @@ -0,0 +1,368 @@ +# 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 +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 tuning, the same at TP4 and TP8). +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, +): + """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. + + 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 + 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), + ) + 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, +): + """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.""" + 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_103__tp8/kernels/m3_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/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_103__tp8/kernels/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_103__tp8/kernels/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/m3_moe_kernel.py new file mode 100644 index 000000000000..a45f0cee2ea3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/m3_moe_kernel.py @@ -0,0 +1,3817 @@ +# 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 +# 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 +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 +# 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 # 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 # 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 + 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 +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 +# 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 +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 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_103__tp8/kernels/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/op.py new file mode 100644 index 000000000000..d9d3fe220d8d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/op.py @@ -0,0 +1,677 @@ +# 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 +INTER = 384 # the routed intermediate per TP rank +NUM_EXPERTS = 128 +SHARED_INTER = 384 # the shared expert's intermediate per TP rank (the routed one) +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) +# ``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 +# 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) + 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, + ): + """``inter``: the routed intermediate per rank; this engine is built for INTER (384) only.""" + 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 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 * 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_103__tp8/kernels/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/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_103__tp8/kernels/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_103__tp8/kernels/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_router_gemv.py new file mode 100644 index 000000000000..88a475f509d5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/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 + + +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_103__tp8/kernels/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_sampled_accept.py new file mode 100644 index 000000000000..315e80947213 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_sampled_accept.py @@ -0,0 +1,158 @@ +# 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 dispatch + +# 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 + + +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_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(): + dispatch.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] > 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) + 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_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), + ) + return accepted, num_accepted, gather_ids diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_shard_accept.py new file mode 100644 index 000000000000..e1dafd292d6d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/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 dispatch + +#: Verify rows per call, and exchange rows per call half of the workspace. +SHARD_ROWS = 48 + + +@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(): + dispatch.capture_missed(dispatch.SHARD_ACCEPT.path, 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_103__tp8/kernels/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/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_103__tp8/kernels/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_103__tp8/kernels/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/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_103__tp8/kernels/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_103__tp8/kernels/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/__init__.py new file mode 100644 index 000000000000..7d2ec0224cf2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/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 TP8 +projection shapes (see op.py).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/m3_mx_gemv_kernel.py new file mode 100644 index 000000000000..b1c9d4de0e08 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/m3_mx_gemv_kernel.py @@ -0,0 +1,1236 @@ +# 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) + + +@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_103__tp8/kernels/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/op.py new file mode 100644 index 000000000000..61c5f6962e2f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/op.py @@ -0,0 +1,472 @@ +# 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 TP8 per-rank projection shapes. + +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[, contig, nacc, a_tmem, st_async]) of the TP8 per-rank projections: GB300 exposed time +# behind a PDL predecessor. The optional fields default to _CONTIG, _NACC, _A_TMEM, _ST_ASYNC. The sparse-layer QKV +# splits K over 8 CTAs (96 CTAs at 1 token; bit-identical to splits 4 and 2). +_CONFIGS: Dict[Tuple[int, int], tuple] = { + # sparse-layer QKV + index projection + (1536, 6144): (8, 8), + # dense-layer QKV + (1280, 6144): (4, 12), + # o_proj + (6144, 1024): (4, 8), + # dense-layer gate_up + (3072, 6144): (4, 12), + # dense-layer down + (6144, 1536): (2, 12), +} +# Per-token-count configs (N, K, M): at 2 tokens the sparse-layer QKV splits K over 4 CTAs, A operand from shared +# memory. +_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 same 4-CTA split for the sparse-layer QKV. +_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) + # 17..32 rows at K < 1024 decline (an unsplit K below 1024 is slower than the FlashInfer GEMM there); no TP8 + # 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 + 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_103__tp8/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/triton_sparse_decode.py new file mode 100644 index 000000000000..fde91fd41c5a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/triton_sparse_decode.py @@ -0,0 +1,1662 @@ +# 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_103 / tp8 target's copy. On top of the +in-tree decode it routes long-context TP8 steps 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 + +# One sparse block is exactly one KV page. +SPARSE_BLOCK_SIZE = 128 + +_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 + +# 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 + + +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 + + 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, + 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 +): + 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): + 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. + page = tl.load(bt_row + blk).to(tl.int64) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len + 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, + 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 +): + 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 + 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 _ in tl.range(0, num_candidates): + if QUERY_GROUP_SIZE == 2: + if chunk_equal: + has_block = True + 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 + 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 + 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. + page = tl.load(bt_row + blk).to(tl.int64) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len_max + q_pos_mask = blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None] + 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 both datacenter Blackwell targets. +_SM100F_CAPABILITIES = ((10, 0), (10, 3)) +# 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 +# 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_lc8: same table, output and MXFP8 hand-off contract.""" + 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_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=linear, + mxfp8=mxfp8, + use_pdl=use_pdl, + ) + if mxfp8 is not None: + m3_prequant.SLOT.publish( + output.view(total_q, num_heads * head_dim), mxfp8[0][:total_q], mxfp8[1] + ) + + +@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; " + 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 is the kernel's (dispatch.SPARSE_DECODE_LC8); operands it declines are a miss. + lc8_covered = ( + kv_nvfp4 + and widen_q + and not identity_blocks + and num_topk_chunks is None + and q.dim() == 3 + and int(q.shape[1]) == m3_attn_lc8.NUM_HEADS + ) + if lc8_covered and not m3_attn_lc8.lc8_applies( + q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output + ): + from .. import dispatch + + dispatch.missed( + dispatch.SPARSE_DECODE_LC8.path, + total_q, + "the step's q / KV pages / scales / output are not the layout the kernel is built for", + ) + lc8_covered = False + if lc8_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) + ) + if kv_nvfp4 and total_q * num_kv_heads <= 64: + # GB300 TP4 Eagle verify (1 KV head, group 16, 6 query tokens): one block per program is fastest up to + # batch 8 (lowm5/bench_decode_chunks.py: b8 11.7 vs 14.8 us with the default split of 4). + num_topk_chunks = 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 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 + ) + 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), + *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, + **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_103__tp8/kernels/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/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_103__tp8/kernels/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_103__tp8/layers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/layers.py new file mode 100644 index 000000000000..f0bdf31bb4c7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/layers.py @@ -0,0 +1,1403 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 NVFP4 / sm_103 / tp8: the decoder layers, on the built-in modules. + +The attention and MoE modules subclass the built-in ``MiniMaxM3Attention`` and +``MiniMaxM3MoE``: the generic path (prefill, mixed steps, wide steps) runs their +code unchanged, and the checkpoint loads through the built-in loader. What this +target adds is the decode-step kernels, each taken only by the step shapes it +was written for and measured on (GB300, tp8 over two trays, 1-8 requests): + +* attention: the NVFP4 QKV + indexer producer at the 8-head geometry; for a + long-context generation step of up to 2 requests the index scoring and the + top-16 block selection in one launch, else the index scorer and the long-row + selector; the selected blocks' attention, split-K merge and the o_proj's MXFP8 + hand-off in one 8-head launch; +* MoE: the router (an fp32 GEMV or a tensor-core router fused with the routed + experts' NVFP4 activation quantization) and the routed experts (width 384) as + one engine launch for 1-8 tokens, with the shared expert folded in; +* the layer boundaries: a multicast all-reduce fused with the residual add and + the next RMSNorm up to 48 tokens. + +The decoder layer and the text model are this target's own: they construct the +subclassed modules. Their generic forward is the built-in one. +""" + +from __future__ import annotations + +from typing import Callable, Dict, Optional, Tuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +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.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +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_minimaxm3 import ( + MiniMaxM3Attention, + MiniMaxM3MoE, + _build_swiglu_oai_dense_mlp, + _nvfp4_cache_supports_fused_write, + _validate_sparse_attention_runtime_config, + get_moe_layer_ids, + get_sparse_disable_index_value_layer_ids, + get_sparse_layer_ids, + maybe_bcg_minimax_m3_attn_custom_op_inplace, +) +from tensorrt_llm._torch.models.modeling_utils import DecoderModel +from tensorrt_llm._torch.modules.decoder_layer import DecoderLayer +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.modules.linear import flashinfer_mxfp8_autotune_active +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph +from tensorrt_llm._torch.speculative import SpecMetadata +from tensorrt_llm._torch.utils import ( + AuxStreamType, + EventType, + Fp4QuantizedTensor, + is_torch_compiling, +) +from tensorrt_llm.logger import logger + +from . import dispatch +from .dispatch import MOE_ENGINE_MAX_TOKENS, PRODUCER_EARLY_TRIGGER, ROUTER_TC_MAX_TOKENS +from .kernels import m3_mc_allreduce, m3_prequant +from .kernels.m3_router_gemv import m3_router_gemv, m3_router_gemv_fp4, m3_router_tc +from .linears import PUSH_ATTN, PUSH_MOE, mark_m3_linear +from .ops.attention.m3_index_score import SPLIT_K, m3_index_score +from .ops.attention.m3_producer import m3_producer +from .ops.moe.m3_moe_engine import m3_moe_engine +from .spec_worker import capture_hidden_states + +# The row ranges of the engines, the router and the fused attention are in dispatch.py. +# The shared expert folds into the routed engine from one token (m3_moe with dedicated shared-expert CTAs). +MOE_FOLD_MIN_TOKENS = 1 +# SMs left to the concurrent shared-expert kernels while a routed engine without the shared fold runs (persistent). +MOE_ENGINE_RESERVED_SMS = 24 +# CTAs per request of the CuTe DSL index scorer at every batch (each block is scored once, so the scores do not depend +# on it), and the request counts up to which it compiles the early-issue schedule (kernels/m3_index_score). +INDEX_SCORE_SPLIT_K = SPLIT_K +INDEX_SCORE_EARLY_MAX_REQUESTS = dispatch.INDEX_SCORE_EARLY_MAX_REQUESTS +# The generic layer's boundaries (prefill, Eagle3 capture layers): the low-latency multicast all-reduce + residual + +# RMSNorm beats the native one-shot all-reduce up to 8 tokens; with the shared + routed add folded in, up to its own +# maximum. +MC_AR_MAX_TOKENS = 8 + +# 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] = {} +# The CuTe DSL index scorer of the sparse decode (its kernels are compiled by prebuild_index_scorer). +INDEX_SCORE_PATH = dispatch.INDEX_SCORE.path + + +def _scorer_runner(): + from .kernels.m3_index_score import IndexScoreRunner + + return IndexScoreRunner + + +def _scorer_key( + dtype: torch.dtype, num_heads: int, query_len: int, head_dim: int, split_k: int, early: bool +) -> tuple: + """The scorer's compile key (``IndexScoreRunner``).""" + return (dtype, int(num_heads), int(query_len), int(head_dim), int(split_k), bool(early)) + + +def scorer_variants(layers, query_lens) -> list: + """The scorer keys a decode step of these layers can take: every sparse layer's (index dtype, index heads per + rank, index dim) at each decode query length, with the early-issue schedule (up to + ``INDEX_SCORE_EARLY_MAX_REQUESTS`` requests, none here) and without it, where the kernel serves the geometry.""" + runner = _scorer_runner() + keys = [] + for layer in layers: + attn = getattr(layer, "self_attn", None) + if ( + attn is None + or not getattr(attn, "is_sparse_attention_layer", False) + or not isinstance(attn.attn, MiniMaxM3MsaSparseAttention) + ): + continue + config = attn.attn.m3_config + dtype = torch.float8_e4m3fn if attn.attn.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) if INDEX_SCORE_EARLY_MAX_REQUESTS else (False,): + key = _scorer_key( + dtype, + config.num_index_heads, + qlen, + config.sparse_index_dim, + INDEX_SCORE_SPLIT_K, + early, + ) + if key not in keys: + keys.append(key) + return keys + + +def prebuild_index_scorer(layers, query_lens, device) -> list: + """Compile every key of ``scorer_variants`` 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 = _scorer_runner() + keys = scorer_variants(layers, 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 + + +_EXPERT_TENSOR_NAMES = ( + "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", +) + + +def _nvfp4_swiglu_oai_experts(experts: nn.Module): + """The routed experts' operands the engines read, or raise if the backend is not TRTLLM-Gen NVFP4 SwiGLU-OAI.""" + backend = getattr(experts, "backend", experts) + tensors = [getattr(backend, name, None) for name in _EXPERT_TENSOR_NAMES] + scale = getattr(backend, "fc31_input_scale", None) + if ( + any(t is None for t in tensors) + or scale is None + or getattr(backend, "scaling_vector_size", 16) != 16 + or getattr(backend, "fc31_act_scale", None) is not None + or not type(backend).quantize_input.__qualname__.startswith("TRTLLMGenNvfp4Base") + ): + raise ValueError("experts backend is not TRTLLM-Gen NVFP4 SwiGLU-OAI") + tensors = [t.data if isinstance(t, nn.Parameter) else t for t in tensors] + scale = scale.data if isinstance(scale, nn.Parameter) else scale + return tensors, scale + + +def _dense_mlp(model_config: ModelConfig, intermediate_size: int, layer_idx: int) -> nn.Module: + """The built-in dense SwiGLU-OAI MLP, its projections on this target's decode GEMM dispatch.""" + mlp = _build_swiglu_oai_dense_mlp( + model_config=model_config, intermediate_size=intermediate_size, layer_idx=layer_idx + ) + mark_m3_linear(mlp.gate_up_proj) + mark_m3_linear(mlp.down_proj) + return mlp + + +class M3MoE(MiniMaxM3MoE): + """MiniMax-M3 MoE with the decode-step router and expert engines. + + Engines are built per (token count, fusion) by ``prebuild_engines`` from + ``build_targets``, before the executor's warmup (building compiles and warms + the kernel); a forward only looks them up, and a step that finds no engine + is a miss and runs the built-in path. + """ + + def __init__( + self, + model_config: ModelConfig, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + layer_idx: int, + ) -> None: + super().__init__(model_config, aux_stream_dict, layer_idx) + self.layer_idx = layer_idx + if self.shared_experts is not None: + mark_m3_linear(self.shared_experts.gate_up_proj) + mark_m3_linear(self.shared_experts.down_proj, push=PUSH_MOE) + # The routed experts' NVFP4 input scale once resolved; False when they cannot take the fused router. + self._fp4_input_scale: Optional[torch.Tensor | bool] = None + self._engines: Dict[Tuple[int, bool, bool], Optional[object]] = {} + + def _router_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """The fp32 router logits: the fp32 GEMV at one row, the tensor-core router up to 48.""" + rows = hidden_states.shape[0] if hidden_states.dim() == 2 else 0 + path = dispatch.router_path(rows) + if path == dispatch.ROUTER_GEMV.path: + logits = m3_router_gemv(hidden_states, self.gate.weight, use_pdl=get_env_enable_pdl()) + if logits is not None: + return logits + dispatch.missed(path, rows, "the router GEMV declined the operands") + elif path == dispatch.ROUTER_TC.path: + fused = m3_router_tc(hidden_states, self.gate.weight, use_pdl=get_env_enable_pdl()) + if fused is not None: + return fused[0] + dispatch.missed(path, rows, "the tensor-core router declined the operands") + else: + dispatch.uncovered(dispatch.ROUTER_TC.path, f"more than {ROUTER_TC_MAX_TOKENS} rows") + return self.gate(hidden_states) + + def _router_fp4(self, hidden_states: torch.Tensor): + """Router logits and the routed experts' NVFP4 activation from one kernel, or None (then + ``_router_logits``).""" + if hidden_states.dim() != 2 or hidden_states.shape[0] > ROUTER_TC_MAX_TOKENS: + return None + if self._fp4_input_scale is None: + try: + _, scale = _nvfp4_swiglu_oai_experts(self.experts) + weight = getattr(self.experts, "backend", self.experts).w3_w1_weight + self._fp4_input_scale = ( + scale if weight.shape[-1] * 2 == hidden_states.shape[-1] else False + ) + except ValueError: + self._fp4_input_scale = False + logger.info( + f"[M3 router FP4] fused router+NVFP4 activation quant: {self._fp4_input_scale is not False}" + ) + rows = hidden_states.shape[0] + path = dispatch.router_path(rows) + if path == dispatch.ROUTER_GEMV.path: + path = dispatch.ROUTER_GEMV_FP4.path + if self._fp4_input_scale is False: + dispatch.missed( + path, + rows, + "the routed experts are not TRTLLM-Gen NVFP4 SwiGLU-OAI on this hidden size", + ) + return None + if path == dispatch.ROUTER_TC.path: + fused = m3_router_tc( + hidden_states, self.gate.weight, self._fp4_input_scale, use_pdl=get_env_enable_pdl() + ) + else: + fused = m3_router_gemv_fp4( + hidden_states, self.gate.weight, self._fp4_input_scale, use_pdl=get_env_enable_pdl() + ) + if fused is None: + dispatch.missed(path, rows, "the router declined the operands") + return fused + + def engine(self, num_tokens: int, shared: bool) -> Optional[object]: + """This layer's prebuilt routed-expert engine for ``num_tokens`` (1..8), router folded in. + + ``shared`` folds the shared expert in as well (routed + shared out). None + when the built-in path must run, after a miss: the engine failed to build, + or ``prebuild_engines`` did not build it (a forward never builds one). + """ + key = (num_tokens, True, shared) + return self._lookup(self._engines, key, dispatch.MOE_ENGINE.path, num_tokens) + + def _lookup(self, engines: dict, key, path: str, num_tokens: int) -> Optional[object]: + if key not in engines: + what = f"layer {self.layer_idx}'s engine" + if torch.cuda.is_current_stream_capturing(): + dispatch.capture_missed(path, num_tokens, what) + else: + dispatch.missed(path, num_tokens, f"{what} was not built before warmup") + return None + engine = engines[key] + if engine is None: + dispatch.missed(path, num_tokens, "its build failed (logged)") + return engine + + def prebuild_engines(self) -> None: + """Build (compile, on first use of each shared state, and warm) every engine a decode step of this layer can + take: ``MOE_ENGINE``'s token counts with the shared expert folded in. Runs from + ``build_targets`` once the weights are real, on the current stream, in the same order on every rank; a build + that fails is logged and its token count falls back (a miss) in the forward. Whether an engine is kept is + agreed across the TP group (one host collective per engine, here): a build that fails on one rank only drops + that engine on every rank, so the ranks never mix an engine (and its boundary push) with the built-in MoE.""" + if self.shared_experts is None: + return + for num_tokens in sorted(dispatch.MOE_ENGINE.keys): + key = (num_tokens, True, True) + if key not in self._engines: + self._engines[key] = self._agreed( + dispatch.MOE_ENGINE.path, + num_tokens, + self._try_build( + dispatch.MOE_ENGINE.path, num_tokens, self._build_engine, shared=True + ), + ) + + def _agreed(self, path: str, num_tokens: int, engine: Optional[object]) -> Optional[object]: + return m3_mc_allreduce.tp_agree_built( + self.mapping, + engine, + f"[M3 dispatch] {path} at {num_tokens} rows: layer {self.layer_idx}'s engine", + ) + + def _try_build(self, path: str, num_tokens: int, build, **kwargs) -> Optional[object]: + try: + return build(num_tokens, **kwargs) + except (ImportError, RuntimeError, ValueError) as err: + logger.warning( + f"[M3 dispatch] {path} at {num_tokens} rows: layer {self.layer_idx}'s engine failed to build: {err}" + ) + return None + + def _build_engine(self, num_tokens: int, shared: bool) -> object: + from .ops.moe import m3_moe_engine as m3_moe_op + + tensors, scale = _nvfp4_swiglu_oai_experts(self.experts) + device = tensors[0].device + sms = torch.cuda.get_device_properties(device).multi_processor_count + kwargs = {} + if shared: + gu, dn = self.shared_experts.gate_up_proj, self.shared_experts.down_proj + kwargs["shared_weights"] = ( + gu.weight.data.view(torch.float8_e4m3fn), + gu.weight_scale.data, + dn.weight.data.view(torch.float8_e4m3fn), + dn.weight_scale.data, + ) + state = _ENGINE_STATES.get((str(device), num_tokens, shared)) + if state is None: + state = _ENGINE_STATES[(str(device), num_tokens, shared)] = m3_moe_op.M3MoeState( + device, + num_tokens=num_tokens, + num_ctas=sms - MOE_ENGINE_RESERVED_SMS, + shared=shared, + ) + engine = state.layer( + *tensors, + self.gate.e_score_correction_bias.data, + router_weight=self.gate.weight.data, + fc31_input_scale=scale.reshape(1).float(), + **kwargs, + ) + # Compile (first use of the state) and warm. + engine( + torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device) + ) + return engine + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + final_all_reduce_params: Optional[AllReduceParams] = None, + lora_params: Optional[dict] = None, + do_finalize: bool = True, + return_parts: bool = False, + ): + """The MoE output; or, with ``return_parts``, ``(shared, routed)`` (routed None when the + engine folded the shared expert in); or, with ``do_finalize=False``, the unreduced + ``[shared_output, fc2_output, expert_scale_factor, expanded_idx_to_permuted_idx]`` + for the fused finalize + all-reduce + residual + RMSNorm epilogue. + + Mirrors the built-in ``MiniMaxM3MoE.forward`` apart from the decode engines. + """ + all_rank_num_tokens = attn_metadata.all_rank_num_tokens + engine_rows = ( + do_finalize + and hidden_states.dim() == 2 + and hidden_states.dtype == torch.bfloat16 + and 1 <= hidden_states.shape[0] <= MOE_ENGINE_MAX_TOKENS + ) + + def _compute_routed_output(): + fused = self._router_fp4(hidden_states) + if fused is not None: + router_logits, x_fp4, x_sf = fused + return self.experts( + Fp4QuantizedTensor(x_fp4, x_sf, is_sf_swizzled=False), + router_logits, + do_finalize=do_finalize, + output_dtype=hidden_states.dtype, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=False, + lora_params=lora_params, + ) + router_logits = self._router_logits(hidden_states) + return self.experts( + hidden_states, + router_logits, + do_finalize=do_finalize, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=False, + lora_params=lora_params, + ) + + def _compute_shared_output(): + return self.shared_experts(hidden_states, lora_params=lora_params) + + folded = None + if ( + engine_rows + and self.shared_experts is not None + and hidden_states.shape[0] >= MOE_FOLD_MIN_TOKENS + ): + engine = self.engine(hidden_states.shape[0], shared=True) + if engine is not None: + arp = m3_prequant.MOE_PUSH.take_engine() + folded = m3_moe_engine(hidden_states.contiguous(), engine, arp=arp) + if arp is not None: + m3_prequant.MOE_PUSH.engine_pushed = folded.data_ptr() + if folded is not None: + # Routed + shared from one kernel. + if return_parts: + return folded, None + result = folded + elif self.shared_experts is None: + result = _compute_routed_output() + else: + routed_output, shared_output = maybe_execute_in_parallel( + _compute_routed_output, + _compute_shared_output, + self.event_dict[EventType.Main], + self.event_dict[EventType.MoeShared], + # 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 self.aux_stream, + disable_on_compile=True, + ) + if not do_finalize: + return [shared_output, *routed_output] + if return_parts: + return shared_output, routed_output + # In-place add into ``shared_output`` to avoid allocating a temporary. + result = shared_output.add_(routed_output) + # tp8 routed experts are tensor-parallel (never fused-communication), so the sum is still rank-local. + return self.allreduce(result, all_reduce_params=final_all_reduce_params) + + +def producer_pw_covered(layer_uses_nvfp4, num_heads_q, num_heads_kv, num_heads_index) -> bool: + """Whether ``dispatch.PRODUCER`` covers a sparse layer: 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 bool(layer_uses_nvfp4) and (num_heads_q, num_heads_kv, num_heads_index) == (8, 1, 1) + + +def _producer_miss(covered: bool, rows: int, reason: str) -> None: + """A producer precondition failed: a miss of ``dispatch.PRODUCER`` on a covered layer (raised under + ``modeling_v2='require'``), otherwise nothing; the caller then runs the unfused producer.""" + if covered: + dispatch.missed(dispatch.PRODUCER.path, rows, reason) + return None + + +def producer_pw_supported(num_heads_q, num_heads_kv, num_heads_index, kv_data, kv_scale) -> bool: + """Whether a sparse layer's NVFP4 producer can take ``m3_producer``: the TP8 heads (8 query, one KV and one + index head) and the NVFP4 pool [pages, 2, 1, 128, 64] with its scale pool [pages, 2, 1, 128, 8].""" + return ( + num_heads_q == 8 + and num_heads_kv == 1 + and num_heads_index == 1 + and 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] + ) + + +class M3Attention(MiniMaxM3Attention): + """MiniMax-M3 attention with the decode-step fused kernels.""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + mark_m3_linear(self.qkv_proj) + mark_m3_linear(self.o_proj, push=PUSH_ATTN) + if not self.is_sparse_attention_layer: + # The KV split trtllm-gen dense decode runs at, per request count (msa_decode.dense_decode_max_seq_len). + self.attn.dense_decode_max_seq_len_by_batch = dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH + + def _check_producer_projection(self) -> None: + """A layer the producer covers (dispatch.PRODUCER) needs the fused QKV/index projection: without it the step is + a miss (raises under require) instead of silently taking the unfused path.""" + if not self.enable_fused_qkv_index_projection and producer_pw_covered( + self.main_kv_is_nvfp4, + self.num_heads, + self.num_key_value_heads, + self.sparse_num_index_heads, + ): + dispatch.missed(dispatch.PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _fused_fp8_qkv_indexer_norm_rope_kv_insert( + self, + packed: torch.Tensor, + position_ids: Optional[torch.Tensor], + attn_metadata: AttentionMetadata, + *, + cache_tensors: Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = None, + ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """Run the vLLM-style horizontal producer for every sparse batch. + + Mirrors the built-in ``MiniMaxM3Attention._fused_fp8_qkv_indexer_norm_rope_kv_insert``; an NVFP4 cache takes + this target's ``m3_producer`` (the native producer bit for bit, its constant loads staged before the grid + wait and its dependents triggered on entry) where its geometry is the one that kernel is built for. + + The CUDA kernel is token-major and batch-type agnostic: per-token + positions and cache slots cover pure prefill, mixed aggregate batches, + and CUDA-graph decode. FP8 caches receive FP8 K/V directly. NVFP4 + caches receive packed E2M1 K/V plus their E4M3 SF16 scale bytes while + Q, index-Q, and index-K remain FP8. + """ + if ( + not self.enable_fused_qkv_index_projection + or not isinstance(self.attn, MiniMaxM3MsaSparseAttention) + or not (self.main_kv_is_fp8 or self.main_kv_is_nvfp4) + or self.attn.indexer_kv_dtype != "fp8" + ): + return None + # A sparse layer on an NVFP4 cache at the TP8 heads is the producer's (dispatch.PRODUCER): every precondition + # below that fails for it is a miss. + covered = producer_pw_covered( + self.main_kv_is_nvfp4, + self.num_heads, + self.num_key_value_heads, + self.sparse_num_index_heads, + ) + rows = int(packed.shape[0]) + rope = self.pos_embd_params.rope if self.pos_embd_params is not None else None + rotary_dim = int(rope.dim) if rope is not None else 0 + if ( + packed.dtype != torch.bfloat16 + or position_ids is None + or self.head_dim != 128 + or self.sparse_index_dim != 128 + or self.sparse_num_index_heads != self.num_key_value_heads + or not self.use_gemma_norm + or self.pos_embd_params is None + or not self.pos_embd_params.is_neox + or self.rotary_emb is None + or rope is None + or rotary_dim != 64 + ): + return _producer_miss( + covered, + rows, + "the packed QKV dtype, position ids, head dims, norm or RoPE configuration", + ) + norm_eps = self.q_norm.variance_epsilon + if any( + module.variance_epsilon != norm_eps + for module in (self.k_norm, self.index_q_norm, self.index_k_norm) + ): + return _producer_miss(covered, rows, "the four RMSNorm epsilons differ") + norm_weights = ( + self.q_norm.weight, + self.k_norm.weight, + self.index_q_norm.weight, + self.index_k_norm.weight, + ) + if any(weight.dtype != torch.bfloat16 or not weight.is_cuda for weight in norm_weights): + return _producer_miss(covered, rows, "a norm weight is not BF16 on the GPU") + + kv_cache_manager = attn_metadata.kv_cache_manager + if cache_tensors is None: + if kv_cache_manager is None: + return _producer_miss(covered, rows, "the step has no KV cache manager") + buffers = kv_cache_manager.get_buffers(self.layer_idx, kv_layout="HND") + index_k_cache = attn_metadata.msa_idx_k_cache(self.layer_idx) + out_cache_loc = getattr(attn_metadata, "msa_out_cache_loc", None) + else: + buffers, index_k_cache, out_cache_loc = cache_tensors + num_tokens = int(packed.shape[0]) + layer_uses_nvfp4 = bool( + getattr(kv_cache_manager, "is_nvfp4_layer", lambda _layer_idx: False)(self.layer_idx) + ) + scale_buffers = ( + kv_cache_manager.get_block_scale_buffers(self.layer_idx, kv_layout="HND") + if layer_uses_nvfp4 + else None + ) + kv_quant_scale = getattr(self.qkv_proj, "inv_kv_scales", None) if layer_uses_nvfp4 else None + if layer_uses_nvfp4: + supports_fused_layout = ( + buffers is not None + and scale_buffers is not None + and _nvfp4_cache_supports_fused_write( + buffers, scale_buffers, self.num_key_value_heads + ) + ) + supported_main_cache = ( + supports_fused_layout + and buffers.is_cuda + and scale_buffers.is_cuda + and scale_buffers.device == buffers.device + and kv_quant_scale is not None + 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_cache = ( + buffers is not None + and buffers.is_cuda + and buffers.dtype == torch.float8_e4m3fn + and buffers.dim() == 5 + and tuple(buffers.shape[1:]) == (2, self.num_key_value_heads, 128, 128) + and buffers.stride(4) == 1 + and buffers.stride(3) == 128 + and buffers.stride(2) == 128 * 128 + and buffers.stride(1) >= self.num_key_value_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_cache = ( + 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] + ) + rotary_cos_sin = self.rotary_emb.rotary_cos_sin + supported_rope_cache = ( + rotary_cos_sin.is_cuda + and rotary_cos_sin.dtype == torch.float32 + and rotary_cos_sin.is_contiguous() + and rotary_cos_sin.dim() == 3 + and tuple(rotary_cos_sin.shape[1:]) == (2, 32) + ) + if ( + not supported_main_cache + or not supported_index_cache + or not supported_rope_cache + 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() < num_tokens + ): + return _producer_miss( + covered and layer_uses_nvfp4, + num_tokens, + "the layer's NVFP4 / index caches, RoPE table, KV quant scale or cache slots are not the " + "layout the producer is built for", + ) + + common_args = ( + out_cache_loc[:num_tokens], + self.num_heads, + self.num_key_value_heads, + self.sparse_num_index_heads, + self.head_dim, + rotary_dim, + norm_eps, + self.q_norm.weight, + self.k_norm.weight, + self.index_q_norm.weight, + self.index_k_norm.weight, + rotary_cos_sin, + position_ids.reshape(-1).contiguous().to(torch.int32), + ) + if layer_uses_nvfp4 and producer_pw_supported( + self.num_heads, + self.num_key_value_heads, + self.sparse_num_index_heads, + buffers, + scale_buffers, + ): + q, index_q = m3_producer( + packed.contiguous(), + buffers.view(torch.uint8), + scale_buffers.view(torch.uint8), + index_k_cache, + common_args[0], + kv_quant_scale, + norm_eps, + *common_args[7:12], + common_args[12], + self.num_heads, + get_env_enable_pdl(), + PRODUCER_EARLY_TRIGGER, + ) + elif layer_uses_nvfp4: + if covered: + dispatch.missed( + dispatch.PRODUCER.path, + num_tokens, + f"the NVFP4 pool {tuple(buffers.shape)} / scale pool " + f"{tuple(scale_buffers.shape)} is not the producer's layout", + ) + q, index_q = torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed.contiguous(), + buffers.view(torch.uint8), + scale_buffers, + index_k_cache, + common_args[0], + kv_quant_scale, + *common_args[1:], + ) + else: + q, index_q = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed.contiguous(), + buffers, + index_k_cache, + *common_args, + ) + return q.flatten(1), index_q.flatten(1) + + def _msa_attention_core(self, q, k, v, idx_q, idx_k, attn_metadata, output): + """The built-in MSA attention core, except a pure generation step of a sparse layer whose caches the + fused producer already wrote: that one runs this target's selection and sparse decode.""" + span = attn_metadata.msa_decode_span + if ( + self.is_sparse_attention_layer + and k is None + and attn_metadata.num_contexts == 0 + and span is not None + and span.row_first == 0 + ): + self._decode_sparse_core(q, idx_q, attn_metadata, output) + return output + return super()._msa_attention_core(q, k, v, idx_q, idx_k, attn_metadata, output) + + def _decode_sparse_core(self, q, idx_q, attn_metadata, output) -> None: + """Top-16 block selection and the block-sparse decode of one pure generation step. + + The selection is the identity when the metadata staged it (every request within 16 blocks). Otherwise a span + of at most ``dispatch.FUSED_SCORE_SELECT_MAX_REQUESTS`` requests scores and selects in one launch; a wider + span runs the CuTe DSL scorer (the per-index-head block maxima, reduced to the KV head) and the long-row + selector (the built-in selector beyond its shapes). + """ + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + MSA_REQUIRED_TOPK, + msa_paged_kv, + select_blocks_from_maxscore, + ) + + from .kernels.long_select import long_select_blocks + from .kernels.triton_sparse_decode import minimax_m3_sparse_attn_decode + + md = attn_metadata + attn = self.attn + config = attn.m3_config + num_tokens = int(q.shape[0]) + rows = int(md.num_seqs) + qlen = int(md.msa_decode_span.query_len) + block_table = md.msa_block_table[:rows] + seq_lens = md.msa_seq_lens_cuda[:rows] + num_kv_heads = int(attn.indexer.config.num_kv_heads) + # The index scorer's kernels are compiled and loaded by build_targets (prebuild_index_scorer), so no step + # compiles one and the identity can be taken from the first decode step. + identity = md._msa_short_ok and md.msa_short_table is not None + fused = None if identity else self._fused_score_select(idx_q, md, rows, qlen, num_tokens) + if identity: + topk = md.msa_short_table[: num_kv_heads * num_tokens * MSA_REQUIRED_TOPK].view( + num_kv_heads, num_tokens, MSA_REQUIRED_TOPK + ) + elif fused is not None: + topk = fused.permute(1, 0, 2) + else: + idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) + idx_k_cache = md.msa_idx_k_cache(self.layer_idx) + max_score = md.msa_proxy_max_score_view( + config.num_index_heads, md.msa_worst_case_max_k_tiles, num_tokens + ) + early = rows <= INDEX_SCORE_EARLY_MAX_REQUESTS + key = _scorer_key( + idx_q_view.dtype, + config.num_index_heads, + qlen, + idx_q_view.shape[-1], + INDEX_SCORE_SPLIT_K, + early, + ) + if key not in _scorer_runner().kernel_cache: + # Outside require mode the op below compiles it (the only scorer this step has). + dispatch.missed( + INDEX_SCORE_PATH, + num_tokens, + f"the index scorer {key} was not compiled before warmup", + ) + 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, + ) + n_valid_blocks = md.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 + ) + 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=num_tokens <= dispatch.SELECT_EARLY_MAX_ROWS, + ) + if table is None: + dispatch.uncovered( + dispatch.LONG_SELECT.path, + f"{max_score_kv.shape[1]} KV blocks, outside {dispatch.LONG_SELECT_MIN_BLOCKS}.." + f"{dispatch.LONG_SELECT_MAX_BLOCKS}: the built-in selector", + ) + table = select_blocks_from_maxscore( + max_score_kv, + topk=MSA_REQUIRED_TOPK, + n_valid_blocks=n_valid_blocks, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + head_major_output=True, + ) + topk = table.permute(1, 0, 2) + manager = md.kv_cache_manager + k_paged, v_paged = msa_paged_kv(manager, self.layer_idx) + nvfp4_args = {} + if manager.is_nvfp4_layer(self.layer_idx): + k_scale, v_scale = _aligned_nvfp4_dequant_scales(attn, self.qkv_proj.kv_scales) + scales = manager.get_block_scale_buffers(self.layer_idx, "HND") + nvfp4_args = dict( + k_block_scale=scales[:, 0], + v_block_scale=scales[:, 1], + k_global_scale=k_scale, + v_global_scale=v_scale, + ) + minimax_m3_sparse_attn_decode( + q.view(num_tokens, self.num_heads, self.head_dim), + k_paged, + v_paged, + topk, + block_table, + seq_lens, + sm_scale=(self.head_dim**-0.5) / float(attn.q_scaling), + output=output.view(num_tokens, self.num_heads, self.head_dim), + decode_query_len=qlen, + identity_blocks=identity, + **nvfp4_args, + ) + + def _fused_score_select(self, idx_q, md, rows: int, qlen: int, num_tokens: int): + """The top-16 table [total_q, num_kv_heads, 16] of a covered decode span from the fused score + select, or None + (after a miss when the span is covered) for the scorer + selector.""" + from .kernels import m3_fused_score_select as fss + + path = dispatch.FUSED_SCORE_SELECT.path + if not dispatch.FUSED_SCORE_SELECT.covers(rows): + dispatch.uncovered( + path, + f"more than {dispatch.FUSED_SCORE_SELECT_MAX_REQUESTS} requests: scorer + selector", + ) + return None + config = self.attn.m3_config + idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) + idx_k_pages = md.msa_idx_k_cache(self.layer_idx).squeeze(1) + max_tiles = int(md.msa_worst_case_max_k_tiles) + n_valid = md.msa_n_valid_blocks + reason = fss.supported(idx_q_view, idx_k_pages, qlen, config.num_kv_heads, max_tiles) + if reason is None and ( + num_tokens != rows * qlen + or n_valid is None + or n_valid.dtype != torch.int32 + or int(n_valid.shape[0]) < num_tokens + ): + reason = "the step's query rows or valid-block counts are not one decode span" + if reason is not None: + dispatch.missed(path, rows, reason) + return None + bufs = fss.buffers(idx_q.device, max_tiles) + if bufs is None: + dispatch.capture_missed(path, rows, "the fused score + select scratch") + return None + return fss.fused_score_select( + idx_q_view, + idx_k_pages, + md.msa_block_table[:rows], + md.msa_seq_lens_cuda[:rows], + n_valid[:num_tokens].contiguous(), + decode_query_len=qlen, + num_kv_heads=config.num_kv_heads, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + max_tiles=max_tiles, + scratch=bufs[0], + counters=bufs[1], + ) + + def _dense_decode_fp8_q(self, attn_metadata) -> bool: + """On decode-only steps of a hybrid NVFP4 cache, the dense layers (whose KV pool is FP8) take FP8 q/k/v + straight from the fused norm/RoPE kernel instead of casting them separately.""" + return ( + self.main_kv_is_nvfp4 + and not self.is_sparse_attention_layer + and self._msa_backend_active() + and attn_metadata.num_contexts == 0 + ) + + def _dense_forward( + self, + position_ids: torch.IntTensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + all_reduce_params: Optional[AllReduceParams] = None, + **kwargs, + ) -> torch.Tensor: + """The built-in dense attention (layers 0-2), FP8 q/k/v on decode steps. + + Mirrors the built-in ``MiniMaxM3Attention._dense_forward``. + """ + if attn_metadata is None: + raise RuntimeError( + f"MiniMax-M3 dense forward (layer {self.layer_idx}) requires attn_metadata; received None." + ) + qkv = self.qkv_proj(hidden_states) + fused_qkv = self._fused_qk_norm_rope( + qkv, + position_ids, + num_heads_q=self.num_heads, + num_heads_k=self.num_key_value_heads, + num_heads_v=self.num_key_value_heads, + head_dim=self.head_dim, + q_norm=self.q_norm, + k_norm=self.k_norm, + out_fp8=self._emit_fp8_main_qkv() or self._dense_decode_fp8_q(attn_metadata), + ) + if fused_qkv is not None: + q, k, v = self._split_main_qkv(fused_qkv) + else: + assert not self._expect_fused_qk_norm_rope(position_ids), ( + f"MiniMax-M3 dense attention (layer {self.layer_idx}) expected the " + f"fused QK-norm+RoPE kernel (bf16 activations, head_dim=" + f"{self.head_dim}) but fell back to the separate path; qkv dtype " + f"is {qkv.dtype} (expected {self.attn_activation_dtype})." + ) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.apply_qk_norm(q, k) + if self.rotary_emb is not None and position_ids is not None: + q, k = self.rotary_emb(position_ids, [q, k]) + o = self._forward_attention_core(q, k, v, None, None, attn_metadata) + return self.o_proj(o, all_reduce_params=all_reduce_params) + + def _sparse_forward( + self, + position_ids: torch.IntTensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + all_reduce_params: Optional[AllReduceParams] = None, + **kwargs, + ) -> torch.Tensor: + """The built-in sparse attention, with the fused decode kernels on the steps they take. + + Mirrors the built-in ``MiniMaxM3Attention._sparse_forward``. + """ + if attn_metadata is None: + raise RuntimeError( + f"MiniMax-M3 sparse forward (layer {self.layer_idx}) requires attn_metadata; received None." + ) + packed_qkv = None + packed_idx_qk = None + self._check_producer_projection() + if self.enable_fused_qkv_index_projection: + if ( + self.register_to_config + and is_torch_compiling() + and isinstance(self.attn, MiniMaxM3MsaSparseAttention) + and self._emit_fp8_main_qkv() + and self.attn.indexer_kv_dtype == "fp8" + ): + kv_cache, index_k_cache = attn_metadata.msa_layer_cache_tensors[self.layer_idx] + q, idx_q = torch.ops.trtllm.minimax_m3_fused_sparse_qkv_producer( + hidden_states, + position_ids, + kv_cache, + index_k_cache, + attn_metadata.msa_out_cache_loc, + self.layer_idx_str, + ) + o = self._forward_attention_core(q, None, None, idx_q, None, attn_metadata) + return self.o_proj(o, all_reduce_params=all_reduce_params) + if self.register_to_config and (is_torch_compiling() or is_in_breakable_cuda_graph()): + packed = torch.ops.trtllm.minimax_m3_qkv_index_proj( + hidden_states, position_ids, self.layer_idx_str + ) + num_tokens = ( + position_ids.shape[-1] if position_ids is not None else hidden_states.shape[0] + ) + output = packed.new_empty( + (num_tokens, self.num_heads * self.head_dim), dtype=self.attn_activation_dtype + ) + maybe_bcg_minimax_m3_attn_custom_op_inplace( + None, None, None, None, None, packed, position_ids, self.layer_idx_str, output + ) + return self.o_proj(output, all_reduce_params=all_reduce_params) + packed = self.qkv_proj(hidden_states) + horizontal = self._fused_fp8_qkv_indexer_norm_rope_kv_insert( + packed, position_ids, attn_metadata + ) + if horizontal is not None: + q, idx_q = horizontal + o = self._forward_attention_core(q, None, None, idx_q, None, attn_metadata) + return self.o_proj(o, all_reduce_params=all_reduce_params) + main_size = self.q_size + 2 * self.kv_size + packed_qkv, packed_idx_qk = packed.split( + [main_size, self.index_q_size + self.index_k_size], dim=-1 + ) + + # The compatibility/fallback main and index branches, overlapped on the aux stream. + def _main_norm_rope(): + qkv = packed_qkv if packed_qkv is not None else self.qkv_proj(hidden_states) + fused_qkv = self._fused_qk_norm_rope( + qkv, + position_ids, + num_heads_q=self.num_heads, + num_heads_k=self.num_key_value_heads, + num_heads_v=self.num_key_value_heads, + head_dim=self.head_dim, + q_norm=self.q_norm, + k_norm=self.k_norm, + out_fp8=self._emit_fp8_main_qkv(), + ) + if fused_qkv is not None: + return self._split_main_qkv(fused_qkv) + assert not self._expect_fused_qk_norm_rope(position_ids), ( + f"MiniMax-M3 sparse attention (layer {self.layer_idx}) expected the fused QK-norm+RoPE kernel" + ) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.apply_qk_norm(q, k) + if self.rotary_emb is not None and position_ids is not None: + q, k = self.rotary_emb(position_ids, [q, k]) + return q, k, v + + def _index_norm_rope(): + idx_qk = ( + packed_idx_qk if packed_idx_qk is not None else self.index_qk_proj(hidden_states) + ) + graph_fp8_indexer = ( + self.register_to_config + and is_torch_compiling() + and isinstance(self.attn, MiniMaxM3MsaSparseAttention) + and self.attn.indexer_kv_dtype == "fp8" + ) + if not graph_fp8_indexer: + fp8_idx_q = self._fused_fp8_index_qk_norm_rope(idx_qk, position_ids, attn_metadata) + if fp8_idx_q is not None: + return fp8_idx_q, None + fused_idx = self._fused_qk_norm_rope( + idx_qk, + position_ids, + num_heads_q=self.sparse_num_index_heads, + num_heads_k=1, + num_heads_v=0, + head_dim=self.sparse_index_dim, + q_norm=self.index_q_norm, + k_norm=self.index_k_norm, + out_fp8=graph_fp8_indexer, + ) + if fused_idx is not None: + return self._split_index_qk(fused_idx) + assert not self._expect_fused_qk_norm_rope(position_ids), ( + f"MiniMax-M3 sparse index branch (layer {self.layer_idx}) expected the fused QK-norm+RoPE kernel" + ) + idx_q, idx_k = idx_qk.split([self.index_q_size, self.index_k_size], dim=-1) + idx_q, idx_k = self.apply_index_qk_norm(idx_q, idx_k) + if self.rotary_emb is not None and position_ids is not None: + idx_q, idx_k = self.rotary_emb(position_ids, [idx_q, idx_k]) + return idx_q, idx_k + + (q, k, v), (idx_q, idx_k) = maybe_execute_in_parallel( + _main_norm_rope, + _index_norm_rope, + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + disable_on_compile=True, + ) + o = self._forward_attention_core(q, k, v, idx_q, idx_k, attn_metadata) + return self.o_proj(o, all_reduce_params=all_reduce_params) + + +def _plain_capture(spec_metadata, layer_id: int, hidden_states, residual) -> None: + """The in-tree Eagle3 hidden-state capture (no fused fc_norm, no early fc fork).""" + capture_hidden_states(spec_metadata, layer_id, hidden_states, residual, None) + + +class M3DecoderLayer(DecoderLayer): + """One MiniMax-M3 transformer block on this target's modules. + + Mirrors the built-in ``MiniMaxM3DecoderLayer``; the layer boundaries add the + multicast all-reduce + residual + RMSNorm up to its token limits. + """ + + def __init__( + self, + model_config: ModelConfig[PretrainedConfig], + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ) -> None: + super().__init__() + config = model_config.pretrained_config + self.hidden_size = config.hidden_size + self.layer_idx = layer_idx + self.mapping = model_config.mapping + self.enable_attention_dp = self.mapping.enable_attention_dp + + _, sparse_layer_ids = get_sparse_layer_ids(config) + disable_index_value_ids = set(get_sparse_disable_index_value_layer_ids(config)) + self.self_attn = M3Attention( + model_config=model_config, + layer_idx=layer_idx, + is_sparse_attention_layer=layer_idx in sparse_layer_ids, + disable_index_value=layer_idx in disable_index_value_ids, + aux_stream=aux_stream_dict[AuxStreamType.Attention], + ) + + _, moe_layer_ids = get_moe_layer_ids(config) + if layer_idx in moe_layer_ids: + self.block_sparse_moe = M3MoE( + model_config=model_config, aux_stream_dict=aux_stream_dict, layer_idx=layer_idx + ) + self.mlp = None + else: + dense_intermediate = int( + getattr(config, "dense_intermediate_size", config.intermediate_size) + ) + self.mlp = _dense_mlp(model_config, dense_intermediate, layer_idx) + self.block_sparse_moe = None + + # Plain (non-Gemma) boundary norms, so they can drive the fused all-reduce + # + residual + RMSNorm epilogue; the loader folds Gemma's (1 + weight) in. + self.input_layernorm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + use_gemma=False, + ) + self.post_attention_layernorm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + use_gemma=False, + ) + + # tp8 without attention DP: every boundary has a cross-rank reduction to fold. + self.pre_feed_forward_fusion = True + self.post_feed_forward_fusion = True + self.allreduce = AllReduce( + mapping=model_config.mapping, + strategy=model_config.allreduce_strategy, + dtype=config.torch_dtype, + ) + # Wired after the weight load to the next layer's input_layernorm (or the final norm). + self.next_layer_layernorm: Optional[RMSNorm] = None + self.next_is_decoder_layer = False + # The Eagle3 hidden-state capture of this layer; the target shell points it at its worker's (fused fc_norm, + # early fc fork). + self.capture: Callable[..., None] = _plain_capture + + def forward( + self, + position_ids: torch.IntTensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + residual: Optional[torch.Tensor], + spec_metadata: Optional[SpecMetadata] = None, + lora_params: Optional[dict] = None, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + all_reduce_params=AllReduceParams(enable_allreduce=False), + lora_params=lora_params, + **kwargs, + ) + capture = spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx) + if self.block_sparse_moe is not None: + return self.forward_moe( + hidden_states, attn_metadata, residual, capture, spec_metadata, lora_params + ) + return self.forward_mlp(hidden_states, residual, capture, spec_metadata, lora_params) + + def _can_mc(self, hidden_states: torch.Tensor, folded: bool = False) -> bool: + """Whether the multicast all-reduce + residual + RMSNorm takes this boundary.""" + return ( + isinstance(hidden_states, torch.Tensor) + and hidden_states.dim() == 2 + and hidden_states.shape[0] + <= (m3_mc_allreduce.MAX_TOKENS if folded else MC_AR_MAX_TOKENS) + and hidden_states.dtype == torch.bfloat16 + and hidden_states.is_contiguous() + ) + + def mc_allreduce_norm(self, x, x2, residual, norm): + """The multicast all-reduce + optional second addend + residual + RMSNorm, or None without a workspace.""" + workspace = m3_mc_allreduce.get_workspace(self.mapping, x.shape[-1]) + if workspace is None: + return None + return m3_mc_allreduce.m3_mc_allreduce_norm( + x, + residual.contiguous(), + norm.weight, + norm.variance_epsilon, + workspace, + use_pdl=get_env_enable_pdl(), + x2=x2, + ) + + def native_allreduce_norm(self, hidden_states, residual, norm): + """The built-in fused all-reduce + residual + RMSNorm.""" + return self.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=norm.weight, + eps=norm.variance_epsilon, + trigger_completion_at_end=False, + ), + ) + + def _pre_feed_forward_norm(self, hidden_states, residual): + if self._can_mc(hidden_states): + fused = self.mc_allreduce_norm( + hidden_states, None, residual, self.post_attention_layernorm + ) + if fused is not None: + return fused + return self.native_allreduce_norm(hidden_states, residual, self.post_attention_layernorm) + + def _next_layer_norm(self, hidden_states, residual, fused_boundary: bool): + """The next layer's input norm at the boundary; ``fused_boundary`` False keeps the module's own reduction.""" + if self.next_layer_layernorm is None: + return hidden_states, residual + parts = hidden_states if isinstance(hidden_states, tuple) else (hidden_states, None) + if not fused_boundary: + return self.next_layer_layernorm(parts[0], residual) + if self._can_mc(parts[0], folded=parts[1] is not None): + fused = self.mc_allreduce_norm(parts[0], parts[1], residual, self.next_layer_layernorm) + if fused is not None: + return fused + hidden_states = parts[0] if parts[1] is None else parts[0].add_(parts[1]) + return self.native_allreduce_norm(hidden_states, residual, self.next_layer_layernorm) + + def _fused_feed_forward_boundary(self, capture: bool) -> bool: + """Whether the feed-forward output goes unreduced to the fused boundary (all-reduce + residual + next norm in + one kernel). An Eagle3 capture layer takes it too and captures the boundary's reduced residual sum; only the + last layer (no next norm, its raw output captured) keeps the module's own all-reduce.""" + return not capture or self.next_layer_layernorm is not None + + def _feed_forward_boundary( + self, hidden_states, residual, fused_boundary, capture, spec_metadata + ): + """The layer's exit: the next norm at the boundary, with this layer's Eagle3 capture.""" + if not capture: + return self._next_layer_norm(hidden_states, residual, fused_boundary) + if not fused_boundary: + self.capture(spec_metadata, self.layer_idx, hidden_states, residual) + return self._next_layer_norm(hidden_states, residual, False) + normed, updated = self._next_layer_norm(hidden_states, residual, True) + self.capture(spec_metadata, self.layer_idx, updated, None) + return normed, updated + + def forward_moe( + self, hidden_states, attn_metadata, residual, capture, spec_metadata, lora_params + ): + hidden_states, residual = self._pre_feed_forward_norm(hidden_states, residual) + fused_boundary = self._fused_feed_forward_boundary(capture) + hidden_states = self.block_sparse_moe( + hidden_states, + attn_metadata, + final_all_reduce_params=AllReduceParams(enable_allreduce=False) + if fused_boundary + else None, + lora_params=lora_params, + return_parts=fused_boundary and self._can_mc(hidden_states, folded=True), + ) + return self._feed_forward_boundary( + hidden_states, residual, fused_boundary, capture, spec_metadata + ) + + def forward_mlp(self, hidden_states, residual, capture, spec_metadata, lora_params): + hidden_states, residual = self._pre_feed_forward_norm(hidden_states, residual) + fused_boundary = self._fused_feed_forward_boundary(capture) + hidden_states = self.mlp( + hidden_states, + final_all_reduce_params=AllReduceParams(enable_allreduce=False) + if fused_boundary + else None, + lora_params=lora_params, + ) + return self._feed_forward_boundary( + hidden_states, residual, fused_boundary, capture, spec_metadata + ) + + +class M3TextModel(DecoderModel): + """The MiniMax-M3 text decoder on this target's layers. + + Its forward belongs to the core in ``modeling.py``, which dispatches each + step to the prefill or the decode target. + """ + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + _validate_sparse_attention_runtime_config(model_config) + super().__init__(model_config) + config = model_config.pretrained_config + self.vocab_size = config.vocab_size + self.aux_stream_dict = { + AuxStreamType.Attention: torch.cuda.Stream(), + AuxStreamType.MoeShared: torch.cuda.Stream(), + AuxStreamType.MoeChunkingOverlap: torch.cuda.Stream(), + } + self.embed_tokens = Embedding( + config.vocab_size, + config.hidden_size, + dtype=config.torch_dtype, + enable_torch_compile_for_embedding=model_config.enable_torch_compile_for_embedding, + ) + self.layers = nn.ModuleList( + [ + M3DecoderLayer(model_config, layer_idx, self.aux_stream_dict) + for layer_idx in range(config.num_hidden_layers) + ] + ) + # Lets the executor enable the automatic FlashInfer MXFP8 path for the decode graphs. + self._use_flashinfer_mxfp8_decode_graph_default = True + self.norm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + use_gemma=False, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/linears.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/linears.py new file mode 100644 index 000000000000..07c0e668c866 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/linears.py @@ -0,0 +1,232 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-graph MXFP8 GEMM dispatch of this target's projections. + +Every MXFP8 projection of the text model (QKV, o_proj, the dense MLPs and the +shared experts) runs the built-in ``MXFP8LinearMethod`` outside decode CUDA +graphs. Inside them (``flashinfer_mxfp8_decode_graph_capture_active``) the +projection's (N, K) and the step's rows pick one path (``dispatch.linear_path``, +row ranges per projection in ``dispatch.py``): + +1. an activation its producer already quantized to MXFP8 (the boundary + all-reduce or the attention merge publish it in ``m3_prequant.SLOT``): the + tcgen05 decode GEMV on the rows it serves, else the decode-graph MXFP8 GEMM + (AutoTuner-selected backend per token bucket, the built-in method's GEMM); +2. a bf16 activation on the rows the GEMV serves quantizing in its prologue; +3. otherwise the built-in method (MXFP8 quantize + GEMM, AutoTuner-selected per bucket). + +No path runs FlashInfer's own per-rank config sweep: its tactics are not merged +across ranks and it profiled wherever the projection ran (the MoE aux stream for +the shared expert), one CUDA graph per tactic. + +A kernel declining rows its path covers is a miss (``dispatch.missed``); the +chain then continues down the list. The GEMVs push their rows straight into +the boundary all-reduce the decode forward armed for the projection's push +slot (``m3_push``, fixed at construction). +""" + +from __future__ import annotations + +from typing import List, Optional, Tuple + +import torch +from torch import nn + +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_gemm_autotuned +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.modules.linear import ( + Linear, + MXFP8LinearMethod, + flashinfer_mxfp8_autotune_active, + flashinfer_mxfp8_decode_graph_capture_active, +) +from tensorrt_llm._torch.utils import is_torch_compiling + +from . import dispatch +from .kernels import m3_mc_allreduce, m3_prequant +from .kernels.mx_gemv import op as mx_gemv +from .ops.gemm.m3_mx_gemv import m3_mx_gemv +from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 + +# The boundary all-reduce a projection's GEMV pushes its rows into, when the decode forward arms one: the attention +# boundary (o_proj) or the MoE boundary's shared-expert set (the shared expert's down projection). +PUSH_ATTN = "attn" +PUSH_MOE = "moe" + + +def mark_m3_linear(module: Linear, push: Optional[str] = None) -> None: + """Mark a projection for this target's decode dispatch; ``adopt_m3_linears`` applies it after the load. + + ``push``: ``PUSH_ATTN`` / ``PUSH_MOE`` when its GEMV pushes into that boundary's all-reduce whenever the decode + forward armed it. + """ + module.m3_push = push + + +def adopt_m3_linears(root: nn.Module) -> None: + """Put every marked projection under ``root`` on the decode dispatch. + + Runs once the weights are real: the quantization method a projection ends + up with is only final after the model applied its quantization config. + It also compiles every decode GEMV kernel a projection can reach here, + before the executor's warmup: on the current (main) stream, the same shapes + in the same order on every rank, and never inside a forward between + collectives or on the MoE aux stream (``prebuild_m3_gemvs``). + """ + use_pdl = get_env_enable_pdl() + compiled = False + for name, module in root.named_modules(): + if not hasattr(module, "m3_push"): + continue + method = module.quant_method + if not isinstance(method, MXFP8LinearMethod) or not method.use_cutlass: + raise ValueError( + f"{name}: the MiniMax-M3 decode dispatch needs a CUTLASS MXFP8 projection, got {type(method).__name__}" + ) + if type(method) is MXFP8LinearMethod: + method.__class__ = M3MXFP8LinearMethod + mx_gemv.ensure_compiled(module.weight, module.weight_scale, use_pdl) + prebuild_m3_gemvs(module, use_pdl) + compiled = True + if compiled: + torch.cuda.synchronize() + + +def _bucket_tokens(rows: int) -> int: + """The token capacity of the multicast all-reduce bucket a ``rows``-token boundary uses.""" + return next(b for b in m3_mc_allreduce.BUCKETS if rows <= b) + + +def gemv_variants(module: Linear) -> List[Tuple[int, bool, Optional[Tuple[int, int, int]]]]: + """Every (rows, bf16 activation, push (world, bucket tokens, mode)) call the decode dispatch can route to this + projection's GEMV (``M3MXFP8LinearMethod.apply``): 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.""" + n, k = module.weight.shape + push_rows = {PUSH_ATTN: dispatch.ATTN_PUSH.keys, PUSH_MOE: dispatch.MOE_PUSH.keys}.get( + module.m3_push, frozenset() + ) + world = int(module.tp_size) + variants = [] + for rows in range(1, dispatch.DECODE_MAX_TOKENS + 1): + for x_bf16, served in ( + (False, rows in dispatch.MX_GEMV_ROWS.get((n, k), ())), + (True, rows <= dispatch.MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0)), + ): + if not served: + continue + variants.append((rows, x_bf16, None)) + if rows in push_rows: + variants.append( + (rows, x_bf16, (world, _bucket_tokens(rows), dispatch.AR_PUSH_MODE)) + ) + return variants + + +def prebuild_m3_gemvs(module: Linear, use_pdl: bool) -> int: + """Compile every GEMV kernel of ``gemv_variants`` (the push ones included) without launching; returns how many + distinct kernels the projection takes.""" + keys = set() + for rows, x_bf16, push in gemv_variants(module): + key = mx_gemv.precompile( + module.weight, module.weight_scale, rows, x_bf16, use_pdl=use_pdl, ar=push + ) + if key is not None: + keys.add(key) + return len(keys) + + +def _push_target(module: Linear): + """The armed boundary all-reduce this projection pushes its output into, if any.""" + if module.m3_push == PUSH_MOE: + return m3_prequant.MOE_PUSH.take_shared() + if module.m3_push == PUSH_ATTN: + return m3_prequant.AR_PUSH.take_target() + return None + + +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 _note_pushed(output: Optional[torch.Tensor], module: Linear) -> None: + if output is None or not mx_gemv.last_pushed: + return + if module.m3_push == PUSH_MOE: + m3_prequant.MOE_PUSH.shared_pushed = output.data_ptr() + else: + m3_prequant.AR_PUSH.pushed = output.data_ptr() + + +class M3MXFP8LinearMethod(MXFP8LinearMethod): + """``MXFP8LinearMethod`` with the decode-graph dispatch in the module docstring.""" + + def apply(self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor]): + if bias is not None or input.dim() != 2: + return super().apply(module, input, bias) + input = input.contiguous() + weight, weight_scale = module.weight, module.weight_scale + use_pdl = get_env_enable_pdl() + decode_graph = not is_torch_compiling() and flashinfer_mxfp8_decode_graph_capture_active() + if not decode_graph: + return super().apply(module, input, bias) + + rows = input.shape[0] + n, k = weight.shape + prequant = m3_prequant.SLOT.take(input) + if module.dtype != torch.bfloat16 or (prequant is None and input.dtype != torch.bfloat16): + path = ( + dispatch.PREQUANT_GEMM_PATH if prequant is not None else dispatch.GENERIC_GEMM_PATH + ) + else: + path = dispatch.linear_path(n, k, rows, prequant is not None) + output = None + if prequant is not None: + act, act_scale = prequant + if path == dispatch.MX_GEMV_PATH: + output = m3_mx_gemv( + act, act_scale, weight, weight_scale, use_pdl=use_pdl, ar=_push_target(module) + ) + _note_pushed(output, module) + if output is None: + dispatch.missed( + path, + rows, + f"the GEMV declined the [{n}, {k}] projection's MXFP8 activation" + + _uncompiled_note(), + ) + if output is None: + if path == dispatch.PREQUANT_GEMM_PATH: + dispatch.uncovered( + dispatch.MX_GEMV_PATH, + f"the [{n}, {k}] projection at rows the GEMV does not serve", + ) + output = mxfp8_gemm_autotuned( + act, + act_scale, + weight, + weight_scale, + module.dtype, + tune=flashinfer_mxfp8_autotune_active(), + ) + return output + if path == dispatch.MX_GEMV_BF16_PATH: + output = m3_mx_gemv_bf16( + input, weight, weight_scale, use_pdl=use_pdl, ar=_push_target(module) + ) + _note_pushed(output, module) + if output is None: + dispatch.missed( + path, + rows, + f"the GEMV declined the [{n}, {k}] projection's input" + _uncompiled_note(), + ) + if output is None: + if path == dispatch.GENERIC_GEMM_PATH: + dispatch.uncovered( + dispatch.MX_GEMV_BF16_PATH, + f"the [{n}, {k}] projection's bf16 input beyond the decode GEMV's rows", + ) + # Stock's decode-graph dispatch for this shape (mxfp8_quantize_gemm_autotuned): same tactics, same bits. + output = super().apply(module, input, bias) + return output diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py new file mode 100644 index 000000000000..de2214dfddfa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py @@ -0,0 +1,624 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ModelingV2 target: MiniMax-M3 NVFP4 / sm_103 / tp8. + +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) +plus one shared expert. The checkpoint is a modelopt mixed-precision export: +NVFP4 routed experts, MXFP8 projections, bf16 router, embedding and head; the KV +cache is NVFP4 for the sparse layers and FP8 for the dense ones. + +``tp8`` is eight GB300s over two trays (multicast all-reduces over MNNVL), +head-split attention (8 query heads and 1 KV head, duplicated, per rank) and +tensor-parallel experts (every expert on every rank, width 384 per rank), no +attention DP. Speculative decoding, when configured, is one-model Eagle3 with +draft length 3: the in-tree BF16 drafter over the full vocabulary and the BF16 +target head, both at the checkpoint's precision; decode steps of up to 8 +requests verify and draft through fused kernels at that precision +(``spec_worker.py``). The decode kernels are specialized for one request +verifying a draft (4 query rows), the deployment this target was measured on. + +**Each step is dispatched once, on the host, from its shape** (``_select_target``): + +* ``DecodeTarget`` takes a pure generation step of at most 48 tokens (up to 8 + requests, each one token or a draft verify of 1 + 3 + extra rows). It runs + the layers with the decode-step kernels of ``layers.py`` and chooses each + layer boundary's all-reduce by the step's token bucket (``DECODE_PLANS``). + Every fast path it can take, the rows each covers and what a step that cannot + take one does (raise under ``modeling_v2='require'``, else log once and fall + back) are in ``dispatch.py``. +* ``PrefillTarget`` takes every other step (prefill, mixed steps, wider decode + steps) on the built-in modules' generic path. + +**What this target asserts rather than adapts**: SM 10.3, the topology above, +the checkpoint shape and quantization (``routing.py``), the MiniMax-M3 MSA +sparse attention backend with its NVFP4 hybrid KV cache, and text-only input: +the vision tower is not built and a step carrying image input raises. +""" + +from __future__ import annotations + +import time +from typing import Any, Dict, List, Optional + +import torch +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + ModelingV2Mode, + resolved_mode, +) +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.distributed import 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_minimaxm3 import ( + MiniMaxM3ForCausalLM, + get_text_model_config, + is_minimax_m3_vl_config, +) +from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM +from tensorrt_llm._torch.modules.linear import flashinfer_mxfp8_decode_graph_capture_active +from tensorrt_llm._torch.speculative import SpecMetadata +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantAlgo + +from . import dispatch +from . import weights as _weights +from .dispatch import AR_PUSH_MODE, EMBED_FIRSTNORM_MAX_TOKENS, DecodePlan, plan_for +from .kernels import m3_mc_allreduce, m3_prequant +from .kernels.mx_gemv import op as mx_gemv +from .layers import M3DecoderLayer, M3MoE, M3TextModel, prebuild_index_scorer +from .linears import adopt_m3_linears +from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm +from .runtime import make_runtime_hooks +from .spec_worker import M3Eagle3Worker, M3SpecDecForwardMixin + +_SM = (10, 3) + + +class MiniMaxM3Tp8ModelingV2(M3TextModel): + """The text model of this target: owns the layers and dispatches each step to a target.""" + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + super().__init__(model_config) + self._embed_owners: Dict[tuple, DecodeEmbedding] = {} + self._embed_identity: Optional[tuple] = None + self._prefill: Optional[PrefillTarget] = None + self._decode: Optional[DecodeTarget] = None + + def build_targets(self) -> None: + """Build the two targets, once the weights are real.""" + table = self.embed_tokens.weight + norm = self.layers[0].input_layernorm + identity = ( + id(table), + table.data_ptr(), + id(norm.weight), + norm.weight.data_ptr(), + float(norm.variance_epsilon), + ) + if self._embed_owners and identity != self._embed_identity: + raise RuntimeError( + "the embedding or the first norm moved after the decode embedding owners were built" + ) + if not self._embed_owners: + self._embed_owners = { + (rows, dtype): DecodeEmbedding(table, norm, rows, dtype) + for rows in range(1, EMBED_FIRSTNORM_MAX_TOKENS + 1) + for dtype in (torch.int32, torch.int64) + } + self._embed_identity = identity + self.prebuild_kernels() + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + + 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, 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 ``dispatch.missed``). + Returns the seconds each part took and the scorer keys.""" + times: Dict[str, object] = {} + mx_gemv.freeze(False) + start = time.perf_counter() + adopt_m3_linears(self) + times["gemv_s"] = time.perf_counter() - start + start = time.perf_counter() + for layer in self.layers: + moe = layer.block_sparse_moe + if isinstance(moe, M3MoE): + moe.prebuild_engines() + torch.cuda.synchronize() + times["moe_s"] = time.perf_counter() - start + start = time.perf_counter() + times["scorer_keys"] = prebuild_index_scorer( + self.layers, self.decode_query_lens(), self.embed_tokens.weight.device + ) + torch.cuda.synchronize() + times["scorer_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)" + % (times["gemv_s"], times["moe_s"], times["scorer_s"], len(times["scorer_keys"])) + ) + return times + + def _select_target(self, attn_metadata: AttentionMetadata, input_ids: Optional[torch.Tensor]): + """Which target runs this step: host-side counts only, fixed per CUDA-graph capture.""" + if ( + attn_metadata.num_contexts == 0 + and input_ids is not None + and input_ids.ndim == 1 + and dispatch.is_decode_step(input_ids.numel()) + ): + return self._decode + return 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, + **kwargs, + ) -> torch.Tensor: + 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") + target = ( + self._select_target(attn_metadata, input_ids) + if inputs_embeds is None + else self._prefill + ) + return target.forward( + attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, **kwargs + ) + + +class PrefillTarget: + """The general case: context rows, mixed steps and wide generation steps, on the generic layer path.""" + + def __init__(self, core: MiniMaxM3Tp8ModelingV2) -> None: + self.core = core + + def forward( + self, attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, **kwargs + ): + core = self.core + hidden_states = core.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds + residual = None + for layer in core.layers: + hidden_states, residual = layer( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + residual=residual, + spec_metadata=spec_metadata, + **kwargs, + ) + return hidden_states + + +class DecodeTarget: + """A pure generation step of at most 48 tokens. + + 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-row engine, pushing its + sum into the boundary all-reduce; the unfinalized MoE with the finalize in + the all-reduce from 9 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. ``dispatch.py`` lists the row + ranges and what a covered step that cannot take its path does. + """ + + def __init__(self, core: MiniMaxM3Tp8ModelingV2) -> None: + self.core = core + self._moe_push_descs: Dict[tuple, Any] = {} + + def forward( + self, attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, **kwargs + ): + core = self.core + num_tokens = int(input_ids.numel()) + plan = plan_for(num_tokens) + hidden_states, residual = self._embed(input_ids, num_tokens) + for layer in core.layers: + if residual is None or ( + spec_metadata is not None and spec_metadata.is_layer_capture(layer.layer_idx) + ): + hidden_states, residual = layer( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + residual=residual, + spec_metadata=spec_metadata, + **kwargs, + ) + else: + hidden_states, residual = self._layer( + layer, plan, position_ids, hidden_states, residual, attn_metadata, kwargs + ) + return hidden_states + + 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 dispatch.EMBED_FIRSTNORM.covers(num_tokens): + dispatch.uncovered( + dispatch.EMBED_FIRSTNORM.path, f"more than {EMBED_FIRSTNORM_MAX_TOKENS} tokens" + ) + return core.embed_tokens(input_ids), None + owner = core._embed_owners.get((num_tokens, input_ids.dtype)) + if owner is None or not input_ids.is_contiguous(): + dispatch.missed( + dispatch.EMBED_FIRSTNORM.path, + num_tokens, + f"no embedding owner for {input_ids.dtype} ids (built for int32 / int64 by build_targets)" + if owner is None + else "the input ids are not contiguous", + ) + return core.embed_tokens(input_ids), None + return m3_embedding_firstnorm( + input_ids, + owner.table, + owner.norm.weight, + owner.eps, + owner.raw, + owner.output, + owner.use_pdl, + ) + + @staticmethod + def _bucket_flags(layer: M3DecoderLayer, num_tokens: int, hidden: int): + """(bucket workspace, its 8-CTA call counters) of the boundary all-reduce for ``num_tokens``, or None.""" + workspace = m3_mc_allreduce.get_workspace(layer.mapping, 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(): + dispatch.capture_missed(path, num_tokens, "the multicast all-reduce workspace") + else: + dispatch.missed(path, num_tokens, "no multicast all-reduce bucket holds the step") + + def _attn_push_target(self, layer: M3DecoderLayer, num_tokens: int, hidden: 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(layer, num_tokens, hidden) + if found is None: + self._no_bucket(dispatch.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, layer: M3DecoderLayer, num_tokens: int, hidden: 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 + (``M3MoE.prebuild_engines``), so it does.""" + from .ops.moe.m3_moe_engine import ar_push_desc + + found = self._bucket_flags(layer, num_tokens, hidden) + if found is None: + self._no_bucket(dispatch.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(): + dispatch.capture_missed( + dispatch.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, layer: M3DecoderLayer, use_mc: bool, x, residual, norm, x2=None, quant=False + ): + """All-reduce + (second addend) + residual + RMSNorm: the multicast kernel, else the built-in all-reduce.""" + if use_mc: + workspace = m3_mc_allreduce.get_workspace(layer.mapping, x.shape[-1]) + fused = ( + None + if workspace is None + else m3_mc_allreduce.m3_mc_allreduce_norm( + x, + residual.contiguous(), + norm.weight, + norm.variance_epsilon, + workspace, + use_pdl=get_env_enable_pdl(), + x2=x2, + quant=quant, + ) + ) + if fused is not None: + return fused + dispatch.missed( + dispatch.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 layer.native_allreduce_norm(x, residual, norm) + + def _prepushed(self, layer, x, residual, norm, quant, prepushed, x2=None): + workspace = m3_mc_allreduce.get_workspace(layer.mapping, x.shape[-1]) + fused = m3_mc_allreduce.m3_mc_allreduce_norm( + x, + residual.contiguous(), + norm.weight, + norm.variance_epsilon, + workspace, + use_pdl=get_env_enable_pdl(), + x2=x2, + quant=quant, + prepushed=prepushed, + ) + if fused is None: + raise RuntimeError( + f"layer {layer.layer_idx}: rows were pushed into an all-reduce that did not run" + ) + return fused + + def _layer( + self, + layer: M3DecoderLayer, + plan: DecodePlan, + position_ids, + hidden_states, + residual, + attn_metadata, + kwargs, + ): + num_tokens = hidden_states.shape[0] + hidden = hidden_states.shape[-1] + # 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: + dispatch.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 dispatch.ATTN_PUSH.covers(num_tokens) + armed = None + if push: + armed = self._attn_push_target(layer, num_tokens, hidden) + m3_prequant.AR_PUSH.arm(armed) + hidden_states = layer.self_attn( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + all_reduce_params=AllReduceParams(enable_allreduce=False), + **kwargs, + ) + if push: + m3_prequant.AR_PUSH.target = None + if push and m3_prequant.AR_PUSH.take_pushed(hidden_states): + hidden_states, residual = self._prepushed( + layer, hidden_states, residual, layer.post_attention_layernorm, False, AR_PUSH_MODE + ) + else: + if armed is not None: + dispatch.missed( + dispatch.ATTN_PUSH.path, + num_tokens, + "the o_proj did not push its rows into the armed boundary all-reduce", + ) + hidden_states, residual = self._boundary( + layer, plan.attn_mc_ar, hidden_states, residual, layer.post_attention_layernorm + ) + + moe = layer.block_sparse_moe + if moe is None: + out = layer.mlp( + hidden_states, final_all_reduce_params=AllReduceParams(enable_allreduce=False) + ) + return self._boundary(layer, plan.attn_mc_ar, out, residual, layer.next_layer_layernorm) + + norm = layer.next_layer_layernorm + quant = plan.qkv_prequant and layer.next_is_decoder_layer + path = dispatch.moe_path(num_tokens) + moe_push = None + if path == dispatch.MOE_ENGINE.path and moe.shared_experts is not None: + # The engine itself runs inside the MoE forward; looked up (or found missing) here first so a push is + # armed only for an engine that will take it. + engine = moe.engine(num_tokens, shared=True) + if engine is not None and plan.moe_mc_ar and decode_graph: + moe_push = self._moe_push_targets(layer, num_tokens, hidden) + if moe_push is not None: + m3_prequant.MOE_PUSH.arm(*moe_push) + if path == dispatch.MC_FINALIZE.path and plan.moe_mc_ar and moe.shared_experts is not None: + if self._bucket_flags(layer, num_tokens, hidden) is None: + self._no_bucket(dispatch.MC_FINALIZE.path, num_tokens) + else: + shared, fc2, scales, perm = moe( + hidden_states, + attn_metadata, + final_all_reduce_params=AllReduceParams(enable_allreduce=False), + do_finalize=False, + ) + 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.weight, + norm.variance_epsilon, + m3_mc_allreduce.get_workspace(layer.mapping, 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" + ) + return fused + parts = moe( + hidden_states, + attn_metadata, + final_all_reduce_params=AllReduceParams(enable_allreduce=False), + return_parts=True, + ) + shared, routed = parts if isinstance(parts, tuple) else (parts, None) + if moe_push is not None: + m3_prequant.MOE_PUSH.disarm() + if routed is None: + # The engine folded the shared expert in and pushed routed + shared into set A. + folded_pushed = m3_prequant.MOE_PUSH.engine_pushed == shared.data_ptr() + m3_prequant.MOE_PUSH.engine_pushed = m3_prequant.MOE_PUSH.shared_pushed = None + if folded_pushed: + return self._prepushed(layer, shared, residual, norm, quant, 2) + else: + routed_pushed, shared_pushed = m3_prequant.MOE_PUSH.take_pushed(shared, routed) + if routed_pushed and shared_pushed: + return self._prepushed(layer, shared, residual, norm, quant, 3, x2=routed) + if routed_pushed or shared_pushed: + raise RuntimeError("only one MoE part pushed into the boundary all-reduce") + dispatch.missed( + dispatch.MOE_PUSH.path, + num_tokens, + "the MoE did not push its output into the armed boundary all-reduce", + ) + return self._boundary(layer, plan.moe_mc_ar, shared, residual, norm, x2=routed, quant=quant) + + +class ModelingV2MiniMaxM3Nvfp4Sm103Tp8(M3SpecDecForwardMixin, MiniMaxM3ForCausalLM): + """MiniMax-M3 NVFP4 / sm_103 / tp8, text only, with one-model Eagle3 or no speculation. + + A subclass of the built-in text model shell for its loader, its KV-cache and + transceiver preferences and its one-engine speculative plumbing; the text + model it wraps, the forward and the speculative worker are this target's. + """ + + 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 (dispatch.py). + dispatch.set_require(resolved_mode(type(self)) is ModelingV2Mode.REQUIRE) + if is_minimax_m3_vl_config(model_config.pretrained_config): + model_config = get_text_model_config(model_config) + 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 + # The one-engine shell's constructor, past the built-in one that builds the built-in text model. + SpecDecOneEngineForCausalLM.__init__( + self, MiniMaxM3Tp8ModelingV2(model_config), model_config + ) + self.runtime_hooks, self._spec_fold = make_runtime_hooks(self) + if self.spec_worker is not None: + self.spec_worker = M3Eagle3Worker.adopt(self.spec_worker) + worker = self.spec_worker + for layer in self.model.layers: + layer.capture = worker.capture_hidden_states + + def load_weights(self, weights: Dict, *args, **kwargs) -> None: + _weights.load(self, weights, *args, **kwargs) + + def setup_aliases(self) -> None: + """Chain each decoder layer's next_layer_layernorm for the fused boundaries.""" + super().setup_aliases() + layers = self.model.layers + for idx, layer in enumerate(layers): + layer.next_is_decoder_layer = idx < len(layers) - 1 + + def cache_derived_state(self) -> None: + super().cache_derived_state() + self.model.build_targets() + if self.spec_worker is not None: + # The fused draft loop's head copy, from the loaded draft weights, before the executor's warmup. + start = time.perf_counter() + if self.spec_worker.prebuild_draft_head(self.draft_model): + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] fused draft head built before warmup: %.1f s" + % (time.perf_counter() - start) + ) + # The drafter linears' low-M kernels and the drafter attention's KV-split cap, before the warmup. + start = time.perf_counter() + built = self.spec_worker.prebuild_drafter_linears(self.draft_model) + 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(self, attn_metadata: AttentionMetadata, *args, **kwargs): + """The fused-verify one-engine forward (``M3SpecDecForwardMixin``), text only.""" + multimodal_params: List[Any] = kwargs.pop("multimodal_params", None) or [] + if any( + param is not None and getattr(param, "multimodal_data", None) + for param in multimodal_params + ): + raise ValueError( + "this MiniMax-M3 target is text only; the request carries multimodal input" + ) + return super().forward(attn_metadata, *args, **kwargs) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/ops/attention/m3_dense_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_dense_decode.py new file mode 100644 index 000000000000..525de997aaf9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_dense_decode.py @@ -0,0 +1,187 @@ +# 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 (``dispatch.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.trtllm_gen_dense_decode import ( # noqa: E501 + minimax_m3_trtllm_gen_dense_decode, +) + +__all__ = ["m3_dense_decode"] + +#: The TP8 dense layers' geometry: 8 query heads and one KV head of dim 128 per rank, an FP8 KV pool. +HEADS = 8 +KV_HEADS = 1 +HEAD_DIM = 128 +#: The split this target hands trtllm-gen at one request (dispatch.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_103}) + + 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 = q.to(torch.float8_e4m3fn) + 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_103__tp8/ops/attention/m3_fused_score_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_fused_score_select.py new file mode 100644 index 000000000000..0457f456024e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_fused_score_select.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MSA index scoring and top-16 block selection of a pure decode span of at most 2 requests, in one launch.""" + +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_index_score import m3_index_score +from .m3_select_long import m3_select_long + +#: Requests per launch the kernel takes. +MAX_REQUESTS = 2 +#: Query rows per request it takes (decode, or a verify of up to 3 drafts). +MAX_DECODE_QUERY_LEN = 4 +#: (KV head x query row) score rows of one request its stage ring holds. +MAX_ROWS = 8 +#: Bytes of the stage ring the select scratch shares. +RING_BYTES = 12 * 16384 + +# 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) + + +def ring_fits(rows: int, max_tiles: int) -> bool: + """Whether a request's select scratch of ``rows`` score rows over ``max_tiles`` blocks fits the stage ring.""" + return rows * ((max_tiles + 3) // 4 * 4) * 4 + rows * 1024 <= RING_BYTES + + +def scratch_shape(max_tiles: int) -> tuple[int, int, int]: + """The scratch [requests, rows, L] that serves every request count and geometry the op takes.""" + return (MAX_REQUESTS, MAX_ROWS, (max_tiles + 3) // 4 * 4) + + +class M3FusedScoreSelect(OpWrapper): + """Fill ``out`` [total_q, num_kv_heads, 16] with each query row's top-16 KV blocks, straight from the FP8 index Q + and the paged FP8 index-K cache. + + The table is ``m3_select_long`` (ascending ids, -1 tail, ties to the smaller block, the last + ``local_blocks`` valid blocks and the first ``init_blocks`` forced in) applied to the per-KV-head maxima over its + index heads of the ``m3_index_score`` scores, with each row's valid-block count ``n_valid`` clamped to + ``max_tiles``. Grid (splits, requests): each CTA scores a share of its request's blocks through a 12-stage TMA + ring with ``mma.sync`` (the products and the 128-term sums of the CuTe DSL scorer, so the same fp32 scores); the + last CTA of a request (an arrival counter it returns to zero) selects. The request's block-table row and length + are read and its older index-K pages are issued before the grid wait; only the last CTA releases the PDL + dependents. ``scratch`` holds the per-request scores between the two phases and needs no initialization; + ``counters`` must be zero before the first call and stays zero between calls. Returns ``out``. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="Eagle3 verify at one request (the TP8 operating point): 4 rows over 20000 tokens", + spec=dict(batch=1, qlen=4, seq_lens=(20000,), max_tiles=8192, **_M3), + ), + Cell( + why="one request decoding one row over 65536 tokens: 512 blocks across every split", + spec=dict(batch=1, qlen=1, seq_lens=(65536,), max_tiles=8192, **_M3), + ), + Cell( + why="one request of 129 tokens: two blocks, most splits empty", + spec=dict(batch=1, qlen=4, seq_lens=(129,), max_tiles=8192, **_M3), + ), + Cell( + why="one request of 200000 tokens: more blocks than one split holds in its 128 local slots", + spec=dict(batch=1, qlen=4, seq_lens=(200000,), max_tiles=8192, **_M3), + ), + Cell( + why="two requests verifying 4 rows each, of different lengths", + spec=dict(batch=2, qlen=4, seq_lens=(20000, 3000), max_tiles=8192, **_M3), + ), + Cell( + why="two requests, one of them shorter than 16 blocks: its whole row is selected and -1 padded", + spec=dict(batch=2, qlen=4, seq_lens=(1500, 70000), max_tiles=8192, **_M3), + ), + ) + + note = """ + Certified at one index head and one KV head per rank (TP8) with M3's forcing (`init_blocks=0`, + `local_blocks=1`), FP8 index Q and index-K, a 1M-token engine (`max_tiles=8192`). The selection is checked bit + for bit against `reference`; scores are compared through the table only. Rows with more than 1024 blocks tied at + the 16th score are not driven. + """ + + def __call__( + self, + idx_q: torch.Tensor, + idx_k_pages: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + n_valid: torch.Tensor, + out: torch.Tensor, + scratch: torch.Tensor, + counters: torch.Tensor, + decode_query_len: int, + num_kv_heads: int, + init_blocks: int, + local_blocks: int, + max_tiles: int, + use_pdl: bool, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_tp8_fused_score_select( + idx_q, + idx_k_pages, + block_table, + seq_lens, + n_valid, + out, + scratch, + counters, + decode_query_len, + num_kv_heads, + init_blocks, + local_blocks, + max_tiles, + use_pdl, + ) + return out + + def reference( + self, + idx_q: torch.Tensor, + idx_k_pages: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + n_valid: torch.Tensor, + out: torch.Tensor, + scratch: torch.Tensor, + counters: torch.Tensor, + decode_query_len: int, + num_kv_heads: int, + init_blocks: int, + local_blocks: int, + max_tiles: int, + use_pdl: bool, + ) -> torch.Tensor: + """The int32 [total_q, num_kv_heads, 16] table: ``m3_select_long.reference`` over the per-KV-head maxima of + ``m3_index_score.reference``. ``out``, ``scratch``, ``counters`` and ``use_pdl`` only say where and how the op + works.""" + total_q, heads, _ = idx_q.shape + placeholder = torch.empty(heads, total_q, max_tiles, device=idx_q.device) + scores = m3_index_score.reference( + idx_q, idx_k_pages, block_table, seq_lens, placeholder, decode_query_len, 1, False + ) # [heads, total_q, max_tiles] + per_kv = scores.view(num_kv_heads, heads // num_kv_heads, total_q, max_tiles).amax(dim=1) + table = torch.empty(total_q, num_kv_heads, 16, dtype=torch.int32, device=idx_q.device) + return m3_select_long.reference( + per_kv.transpose(1, 2), n_valid, table, init_blocks, local_blocks + ) + + def is_valid( + self, + idx_q: torch.Tensor, + idx_k_pages: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + n_valid: torch.Tensor, + out: torch.Tensor, + scratch: torch.Tensor, + counters: torch.Tensor, + decode_query_len: int, + num_kv_heads: int, + init_blocks: int, + local_blocks: int, + max_tiles: int, + use_pdl: bool, + ) -> None: + """The inputs this op takes and answers wrongly. + + A non-zero counter makes a request's selection run early, on partial scores. The block-table row is read up + to ``ceil(seq_len / 128)`` entries unchecked. Reads ``counters`` and ``seq_lens`` back to the host: a test-only + cost. + """ + assert bool((counters[: seq_lens.shape[0]] == 0).all()), ( + "counters must be zero on entry: the kernel selects once a request's arrivals reach its split count" + ) + nblocks = (seq_lens + 127) // 128 + assert bool((nblocks <= block_table.shape[1]).all()), ( + "a request runs past its block-table row" + ) + assert bool((nblocks <= max_tiles).all()), "a request has more blocks than max_tiles" + assert init_blocks >= 0 and local_blocks >= 0, "forcing counts must be non-negative" + + +m3_fused_score_select = M3FusedScoreSelect() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_index_score.py new file mode 100644 index 000000000000..7c9850fb3310 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_index_score.py @@ -0,0 +1,153 @@ +# 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 ``kernels/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 ...kernels.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_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="three requests verifying 1 + 3 rows (the narrowest span the fused score + select leaves to it)", + spec=dict(seq_lens=(20000, 3000, 129), qlen=4, early=False), + ), + Cell( + why="three requests, one longer than 512 blocks (the CTAs loop)", + spec=dict(seq_lens=(70000, 2049, 4096), 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 + ), + ), + Cell( + why="eight requests decoding one row each", + spec=dict( + seq_lens=(300, 1500, 2049, 4096, 5000, 7777, 8192, 33001), qlen=1, 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_103__tp8/ops/attention/m3_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_producer.py new file mode 100644 index 000000000000..7dcfbd4624b9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_producer.py @@ -0,0 +1,197 @@ +# 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_103}) + + 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, + ) -> tuple[torch.Tensor, torch.Tensor]: + 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, + ) -> 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, + ) -> 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_103__tp8/ops/attention/m3_select_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_select_long.py new file mode 100644 index 000000000000..867822e7fe41 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_select_long.py @@ -0,0 +1,172 @@ +# 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_103}) + + 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 (TP8) 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, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_select_long( + scores, n_valid_blocks, out, init_blocks, local_blocks, use_pdl, early_trigger + ) + 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, + ) -> 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`` and + ``early_trigger`` 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, + ) -> 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_103__tp8/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode.py new file mode 100644 index 000000000000..648c03e76089 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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 ...kernels.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 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; 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_103}) + + 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 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 each route publishes is certified on the long-context + kernel's own entry (`m3_sparse_decode_lc8`); 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_103__tp8/ops/attention/m3_sparse_decode_lc8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode_lc8.py new file mode 100644 index 000000000000..d0676c3495a2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode_lc8.py @@ -0,0 +1,433 @@ +# 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 +) -> 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 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_103}) + + 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, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_tp8_sparse_decode_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, + ) + 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, + ) -> torch.Tensor: + """The BF16 ``out`` [T, 1024] the op is supposed to write: ``sparse_attention_reference`` in float64, + rounded once. ``linear_softmax``, ``use_pdl`` and ``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, + prewait: bool, + ) -> 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" + + +m3_sparse_decode_lc8 = M3SparseDecodeLc8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/msa_kv_update.py new file mode 100644 index 000000000000..ad05b934d4a9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp8/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_mc_allreduce_norm.py new file mode 100644 index 000000000000..2aad4aa72415 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_mc_allreduce_norm.py @@ -0,0 +1,277 @@ +# 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, 8 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, +) -> tuple[torch.Tensor, ...]: + """The boundary arithmetic once each rank's BF16 partial is known (shared with the finalize entry). + + Ranks are summed in FP32 in rank order, the residual added and rounded to BF16 once; 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() + updated = (total + residual.float()).to(torch.bfloat16) + u = 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_103}) + + # The tp8 decode boundaries (hidden 6144, 8 ranks over two trays, 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 (dispatch.MC_ALLREDUCE): the bucket and epilogue options the + # target picks for it (kernels/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): + # dispatch.ATTN_PUSH / dispatch.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) + ), + ) + + 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, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_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, + ) + + 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, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute. + + 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) + + 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, + ) -> 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_103__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..64feae8122d5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/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, 8 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_103}) + + # The tp8 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 (dispatch.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_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_103__tp8/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/ops/gemm/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_draft_lowm.py new file mode 100644 index 000000000000..4c3eca583b61 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_draft_lowm.py @@ -0,0 +1,150 @@ +# 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; rows +#: past the last entry are not served. GB300 cold L2, CUDA-graph timing, us cuBLAS -> low-M at 1 / 4 rows: qkv 11.9 -> +#: 6.4 / 12.0 -> 7.8, o 3.6 -> 2.3 / -, gate_up 11.9 -> 9.4 / 12.1 -> 11.0, down 6.8 -> 5.6 / -, fc 37.9 -> 31.9 / +#: 39.2 -> 35.5. +TACTICS: Mapping[Tuple[int, int], tuple] = { + # qkv + (1280, 12288): ( + (1, "d", (384, 2, 1)), + (2, "d", (128, 4, 2)), + (3, "d", (128, 4, 3)), + (8, "s", (64, 8, 4, 12)), + ), + # o_proj + (6144, 1024): ((1, "d", (32, 4, 1)), (2, "d", (32, 4, 2))), + # gate_up + (4608, 6144): ( + (1, "d", (128, 1, 1)), + (2, "d", (256, 2, 2)), + (3, "s", (128, 8, 3, 6)), + (4, "s", (128, 8, 3, 5)), + (8, "s", (64, 8, 2, 10)), + ), + # down + (6144, 2304): ((1, "d", (96, 1, 1)), (2, "d", (96, 2, 2))), + # the Eagle3 fc over the three captured hidden states + (6144, 18432): ( + (1, "d", (384, 2, 1)), + (2, "d", (64, 4, 2)), + (3, "d", (32, 2, 3)), + (8, "s", (64, 8, 1, 12)), + ), +} + + +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 + + +def _cells() -> tuple[Cell, ...]: + names = { + (1280, 12288): "qkv", + (6144, 1024): "o_proj", + (4608, 6144): "gate_up", + (6144, 2304): "down", + (6144, 18432): "fc", + } + return tuple( + Cell( + why=f"drafter {names[nk]}, {m} rows: {tactic_for(nk[0], nk[1], m)[0]} tactic", + spec=dict(nk=nk, m=m), + ) + for nk, table in TACTICS.items() + for m in range(1, table[-1][0] + 1) + ) + + +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 (kernels/m3_draft_lowm.prebuild). + """ + + ARCHS = frozenset({Arch.SM_103}) + CELLS: tuple[Cell, ...] = _cells() + + 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, intermediate 18432) and the replicated fc. 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, + ) -> 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)) + 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, + ) -> 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, + ) -> 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" + ) + + +m3_draft_lowm = M3DraftLowM() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv.py new file mode 100644 index 000000000000..0aad4c8ef78b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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 ...kernels.mx_gemv import op as _gemv +from ..moe.m3_moe_engine import mxfp8_dequant + +__all__ = ["PROJECTIONS", "gemv_is_valid", "m3_mx_gemv", "served_cells"] + +# (N, K) of this target's MXFP8 projections at tp8, the GEMV's tuned shapes. The shared expert's projections +# (768 x 6144, 6144 x 384) run inside the MoE engine at decode rows and on the generic GEMM otherwise. +QKV_SPARSE = (1536, 6144) +QKV = (1280, 6144) +O_PROJ = (6144, 1024) +DENSE_GATE_UP = (3072, 6144) +DENSE_DOWN = (6144, 1536) +PROJECTIONS = { + "sparse-layer QKV": QKV_SPARSE, + "dense-attention QKV": QKV, + "o_proj": O_PROJ, + "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_103}) + + # 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 sparse-layer QKV takes another split at 2 tokens and above 8. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token (split 8)", spec=dict(nk=QKV_SPARSE, m=1)), + Cell( + why="sparse-layer QKV, 2 tokens: its split-4 config, A from shared memory", + spec=dict(nk=QKV_SPARSE, m=2), + ), + Cell( + why="sparse-layer QKV, 4 rows: one request verifying 3 drafts", + spec=dict(nk=QKV_SPARSE, m=4), + ), + Cell( + why="sparse-layer QKV, 12 rows: the wide split-4 config", + spec=dict(nk=QKV_SPARSE, m=12), + ), + 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 4)", 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="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 (dispatch.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 and M > 48 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_103__tp8/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..bbb0a98521d1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/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 ...kernels.mx_gemv import op as _gemv +from ..moe.m3_moe_engine import mxfp8_quant_dequant +from .m3_mx_gemv import ( + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + 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_103}) + + # 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, 2 tokens: its split-4 config", spec=dict(nk=QKV_SPARSE, m=2)), + 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, 48 rows: the 48-column MMA", spec=dict(nk=O_PROJ, m=48)), + Cell( + why="dense gate_up, 16 rows: the most split 4 at K 6144 serves", + spec=dict(nk=DENSE_GATE_UP, m=16), + ), + Cell(why="dense down, one token", spec=dict(nk=DENSE_DOWN, m=1)), + Cell( + why="dense down, 32 rows: the most split 2 at K 1536 serves", + spec=dict(nk=DENSE_DOWN, m=32), + ), + ) + # Every other (projection, rows) the decode dispatch can send here (dispatch.MX_GEMV_BF16_MAX_TOKENS). + CELLS = CELLS + served_cells(CELLS, x_bf16=True, max_m=48) + + 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_103__tp8/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_moe_engine.py new file mode 100644 index 000000000000..f46935ffe8ca --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_moe_engine.py @@ -0,0 +1,325 @@ +# 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 TP8 build: routed and shared +expert width 384 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.quantization.utils.fp4_utils import ( + get_reorder_rows_for_gated_act_gemm_row_indices, + get_shuffle_matrix_a_row_indices, +) + +from ...kernels.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_103}) + + # 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 TP8 operating point), records 8x, logits 32x", + 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: 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_103__tp8/ops/moe/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_router_gemv.py new file mode 100644 index 000000000000..37870d0e6cd7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + # 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_103__tp8/ops/moe/m3_router_gemv_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_router_gemv_fp4.py new file mode 100644 index 000000000000..58f809e82cf6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + # 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_103__tp8/ops/moe/m3_router_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_router_tc.py new file mode 100644 index 000000000000..e88af51f7fa2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + # 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 (dispatch.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_103__tp8/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/m3_embedding_firstnorm.py new file mode 100644 index 000000000000..ff7fb16e2266 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/m3_embedding_firstnorm.py @@ -0,0 +1,89 @@ +# 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 tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.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_103}) + + # 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 (dispatch.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_103__tp8/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/ops/spec/m3_capture_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_capture_norm.py new file mode 100644 index 000000000000..131269ae7d07 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + # 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_103__tp8/ops/spec/m3_chunk_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_chunk_norm.py new file mode 100644 index 000000000000..f6adb2207a20 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/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_103}) + + # 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_103__tp8/ops/spec/m3_draft_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_draft_input.py new file mode 100644 index 000000000000..75f7a44dfcfb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + # Every pass of the fused draft loop (dispatch.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_103__tp8/ops/spec/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_draft_step.py new file mode 100644 index 000000000000..3f00def536b6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_draft_step.py @@ -0,0 +1,359 @@ +# 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 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 .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 TP8 head shard (200064-id vocabulary) and hidden size. +DRAFT_LEN = 3 +DRAFT_SHARD = 25008 +#: Ranks of the TP group the argmax is exchanged over. +WORLD = 8 +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. The op also streams a 12-bit + coded head (``exc_off`` given); this target never builds one, and nothing here certifies it.""" + if exc_off is not None: + raise NotImplementedError("the tp8 target streams the BF16 fragments only") + return from_fragments(wf) + + +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 ..``), + (the op also takes a 12-bit coded stream with ``exc_off`` / ``exc_pos`` / ``exc_w`` / ``exp_base``, which this + target never builds); 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_103}) + + # Every pass of the fused draft loop (dispatch.DRAFT_STEP): draft_len 3 at 1..8 requests, the drafter's + # full-vocabulary BF16 head (25008 rows per rank of 8, 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 % WORLD, + world=WORLD, + 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 7's ids", + spec=dict( + nreq=2, draft_len=DRAFT_LEN, step=1, rank=7, world=WORLD, 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=WORLD, 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=5, + world=WORLD, + vs=DRAFT_SHARD, + k=HIDDEN, + ties=True, + ), + ), + ) + # The BF16 fragments at every pass shape. + CELLS: tuple[Cell, ...] = tuple( + Cell(why=f"bf16: {c.why}", spec=dict(c.spec, coded=False)) for c in _PASSES + ) + + note = """ + The head is the in-tree BF16 head's arithmetic (BF16 weights and activations, fp32 accumulation, a BF16 logit) + 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, + ) -> None: + torch.ops.trtllm.minimax_m3_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, + ) + + 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, + ) -> 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, + ) -> 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" + 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_103__tp8/ops/spec/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_sampled_accept.py new file mode 100644 index 000000000000..d0ad11de2ce9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_sampled_accept.py @@ -0,0 +1,269 @@ +# 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_103}) + + # draft_len 3 (T = 4 rows per request) at 1 / 2 / 8 requests (4 / 8 / 32 rows, 8 CTAs per row) and the 48-row + # maximum (4 CTAs per row); the 25008-column TP8 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=25008, 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=25008, + vocab_offset=2 * 25008, + 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=25008, + vocab_offset=25008, + 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=25008, + vocab_offset=3 * 25008, + 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_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_103__tp8/ops/spec/m3_sampled_accept_ws_words.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_sampled_accept_ws_words.py new file mode 100644 index 000000000000..b0a17e38a983 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/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_sampled_accept`` exchanges through for ``world`` ranks: + two call halves of ``[attempts][rows][world * splits][record words]``.""" + + ARCHS = frozenset({Arch.SM_103}) + + 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)), + ) + + 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_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_103__tp8/ops/spec/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_shard_accept.py new file mode 100644 index 000000000000..500553c2bc88 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/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_103}) + + # Every greedy verify step of the target (dispatch.SHARD_ACCEPT): draft_len 3 at 1..12 requests (4..48 rows), the + # target lm_head's BF16 shard logits (25008 columns per rank at TP8). 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 8, next_new staged", + spec=dict(nreq=1, draft_len=3, vs=25008, world=8, rank=0, force="none", strided=False), + ), + Cell( + why="cap 2: 8 verify rows, rank 2, global ids offset by 2 x 25008", + spec=dict(nreq=2, draft_len=3, vs=25008, world=8, rank=2, force="none", strided=False), + ), + Cell( + why="cap 4: 16 verify rows (4 requests x T = 4), rank 1 of 8, next_new staged", + spec=dict(nreq=4, draft_len=3, vs=25008, world=8, rank=1, force="none", strided=False), + ), + Cell( + why="cap 8: 32 verify rows, rank 7, global ids offset by 7 x 25008", + spec=dict(nreq=8, draft_len=3, vs=25008, world=8, rank=7, 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=25008, world=8, 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=25008, + world=8, + rank=2, + force="total", + strided=False, + ties=True, + ), + ), + # Every other request count of the verify range, on the 8-rank 25008-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=25008, world=8, rank=n % 8, 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_103__tp8/ops/spec/m3_slot_scatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_slot_scatter.py new file mode 100644 index 000000000000..470a6372f2b6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + # 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_103__tp8/ops/spec/m3_step_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_step_gather.py new file mode 100644 index 000000000000..abea2cf620e9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/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_103}) + + # 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_103__tp8/runtime.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/runtime.py new file mode 100644 index 000000000000..cfcff73e3cea --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/runtime.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The step-loop hooks this target opts into (``RuntimeHooks``), measured on GB300 AgentX decode (tp4 and tp8). + +* ``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. +* 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. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional, Tuple + +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy, RuntimeHooks + +from .kernels import msa_kv_update +from .spec_fold import M3SpecStepFold + +if TYPE_CHECKING: + from .modeling import ModelingV2MiniMaxM3Nvfp4Sm103Tp8 + + +def make_runtime_hooks( + model: "ModelingV2MiniMaxM3Nvfp4Sm103Tp8", +) -> Tuple[RuntimeHooks, Optional[M3SpecStepFold]]: + """This target's hooks, and the speculative step fold they carry (None without speculation).""" + spec_fold = M3SpecStepFold() if model.spec_config is not None else None + hooks = RuntimeHooks( + pack_step_inputs=True, + skip_unchanged_h2d=True, + loop_on_execution_stream=True, + merge_tuned_tactics_across_ranks=True, + skip_greedy_rng_upload=True, + spec_step_fold=spec_fold, + msa=MsaRuntimePolicy( + short_sequence_identity=True, + kv_update=msa_kv_update.msa_kv_update, + skip_unchanged_block_offsets=True, + reuse_pool_views=True, + ), + ) + return hooks, spec_fold diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec_fold.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec_fold.py new file mode 100644 index 000000000000..dbc210c0ab46 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec_fold.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""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 never stores the sampler's outputs itself, so ``begin_step`` and +``sampler_rows_stored`` always report that nothing was staged. + +* ``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``). +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from .kernels import m3_spec_tail + + +class M3SpecStepFold: + """``SpecStepFold`` implementation of the MiniMax-M3 tp8 target (see the module docstring). + + One instance per target model; the target passes it to ``RuntimeHooks.spec_step_fold``. + """ + + def __init__(self) -> None: + self.store: Optional[object] = None + self.step_inputs: Optional[ + tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] + ] = 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) + + def bind_sampler_store(self, store: object) -> None: + self.store = store + + def begin_step(self, seq_slots: Optional[list[int]], num_rows: int) -> bool: + return False + + 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: + return False + + 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, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec_worker.py new file mode 100644 index 000000000000..ce1f5bf51968 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec_worker.py @@ -0,0 +1,1119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Eagle3 one-model speculative decoding of the MiniMax-M3 tp8 target, at stock numerics. + +``M3Eagle3Worker`` is the in-tree ``Eagle3OneModelWorker`` with a fused verify and a fused draft loop: + +* Verify: on decode-only batches whose verify rows the fused kernels cover, the target's BF16 lm_head computes this + rank's vocabulary shard of the logits (the in-tree LM 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 in-tree Eagle3 drafter's modules and weights at their precision (BF16 weights and activations, the + full-vocabulary BF16 head, greedy drafts). The first pass may take the per-layer fc_norm RMSNorms in one launch (or + folded into the hidden-state capture) and run its fc beside the target head GEMM. Its BF16 linears run + FlashInfer's fixed-tactic low-M kernels at the rows ``dispatch.DRAFT_LOWM`` covers (``prebuild_drafter_linears``). + On a decode-only step of at most 8 requests (``dispatch.DRAFT_STEP``) the draft loop is fused + (``_m3_fused_draft_loop``): each pass's draft input in one launch, the drafter layer's in-tree 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. + +The target's outer CausalLM drives the worker through ``M3SpecDecForwardMixin.forward`` (or the same calls made from +its own forward), and its layer loop captures the Eagle3 hidden states through ``capture_hidden_states``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Optional + +import torch +import triton +import triton.language as tl +from torch import nn + +from tensorrt_llm._torch.distributed import AllReduce, AllReduceParams +from tensorrt_llm._torch.models.modeling_speculative import _slice_spec_position_ids +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager +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, +) +from tensorrt_llm.mapping import Mapping + +from . import dispatch +from .kernels import m3_draft_lowm, m3_draft_step, m3_sampled_accept, m3_shard_accept, m3_spec_tail + +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.llmapi.llm_args import EagleDecodingConfig + +# Id of the target's final hidden state in ``layers_to_capture`` (SpecDecOneEngineForCausalLM.layer_idx). +FINAL_HIDDEN_LAYER_ID = -1 + + +@triton.jit +def _m3_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) + + +def draft_ar_norm( + all_reduce: AllReduce, x: torch.Tensor, residual: torch.Tensor, norm: nn.Module +) -> tuple[torch.Tensor, torch.Tensor]: + """One draft-layer boundary in stock Eagle3's structure: the projection's plain all-reduce (BF16 out), then the + norm module with the residual (FlashInfer's fused_add_rmsnorm: FP32 add, norm of the FP32 sum). Returns the + normalized rows and the residual sum.""" + return norm(all_reduce(x), residual) + + +def capture_hidden_states( + spec_metadata: Any, + layer_id: int, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor], + capture_norm: Optional[tuple[torch.Tensor, float]], +) -> 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) (``M3Eagle3Worker.capture_norm``), 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 " + f"buffer capacity: num_tokens={num_tokens}, " + f"capacity={spec_metadata.hidden_states.shape[0]}" + ) + if num_tokens > hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds " + f"available hidden states: num_tokens={num_tokens}, " + f"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 + + +class M3Eagle3Worker(Eagle3OneModelWorker): + """``Eagle3OneModelWorker`` with the MiniMax-M3 fused verify (see the module docstring). + + Methods whose docstring says "Mirrors" are copies of the named in-tree method with the fused-verify branches + inserted; compare them with that method when it changes. + """ + + def __init__( + self, + spec_config: "EagleDecodingConfig", + mapping: Optional[Mapping] = None, + model_config: Optional["ModelConfig"] = None, + use_separate_draft_kv_cache: bool = False, + ): + super().__init__( + spec_config, + mapping=mapping, + model_config=model_config, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + ) + self._init_fused_verify() + + @classmethod + def adopt(cls, worker: Eagle3OneModelWorker) -> "M3Eagle3Worker": + """Turn the in-tree worker the one-engine shell built (``get_spec_worker``) into this worker, in place.""" + if type(worker) is not Eagle3OneModelWorker: + raise TypeError(f"expected an Eagle3OneModelWorker, got {type(worker).__name__}") + worker.__class__ = cls + worker._init_fused_verify() + return worker + + def _init_fused_verify(self) -> None: + # 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_norm[0].weight, stacked fc_norm weights [layers, H]), built eagerly on the first fused first pass. + self._m3_fc_norm_w: Optional[tuple[torch.Tensor, torch.Tensor]] = None + # Armed per step by arm_capture_norm: (stacked fc_norm weights, eps) or None. + self._m3_capture_norm: Optional[tuple[torch.Tensor, float]] = 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 + # (capture_hidden_states), and the fork that did it, joined by _m3_target_accept. + self._m3_fc_early_drafter: Optional[nn.Module] = None + self._m3_fc_early_stream: Optional[torch.cuda.Stream] = None + self._m3_force_vec_warm = False + # The fragment-ordered copy of the drafter's BF16 head shard the fused draft loop streams, built before the + # executor's warmup (prebuild_draft_head). + self._m3_draft_head: Optional[m3_draft_step.DraftHead] = 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 + + @property + def capture_norm(self) -> Optional[tuple[torch.Tensor, float]]: + """What ``capture_hidden_states`` needs for this step's captures (set by ``arm_capture_norm``).""" + return self._m3_capture_norm + + def capture_hidden_states( + self, + spec_metadata: Any, + layer_id: int, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor] = None, + ) -> None: + """``capture_hidden_states`` with this step's capture norm; 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``).""" + capture_hidden_states( + 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 _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 + ) + ) + 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 because + # the accepted token count affects which Mamba states are valid. The + # isinstance check below naturally no-ops on non-Mamba kv_cache_managers, + # so this is safe to run unconditionally regardless of spec mode (Eagle3 + # over a Mamba-style draft is plausible, even if no such draft exists today). + 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: torch.LongTensor, + position_ids: torch.LongTensor, + hidden_states: torch.Tensor, + accepted_tokens: torch.Tensor, + attn_metadata: Any, + spec_metadata: Any, + draft_model: nn.Module, + ): + """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: Any, spec_metadata: Any) -> 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) > dispatch.DECODE_MAX_TOKENS: + return False + return not spec_metadata.wants_advanced_draft_sampling + + def tail_eligible(self, attn_metadata: Any, spec_metadata: Any, 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 ``dispatch.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 self._m3_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 dispatch.SHARD_ACCEPT.covers(key): + dispatch.uncovered( + dispatch.SHARD_ACCEPT.path, + f"draft length {draft_len} 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: Any, spec_metadata: Any, draft_model: Optional[nn.Module] + ) -> 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.""" + return self._m3_structural_ok(attn_metadata, spec_metadata) and self._m3_fc_norm_ok( + draft_model + ) + + def _m3_seed_first_pass( + self, num_accepted_tokens: torch.Tensor, spec_metadata: Any, 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: torch.Tensor, + accepted_tokens: torch.Tensor, + num_tokens: int, + num_contexts: int, + spec_metadata: Any, + draft_model: nn.Module, + ) -> Optional[dict[str, Any]]: + """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 in-tree 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 not self._m3_fc_norm_ok(draft_model) + ): + return None + dm = draft_model.model + layers = len(spec_metadata.layers_to_capture) + if dm.hidden_size_in != spec_metadata.hidden_size or layers != len(dm.fc_norm): + return None + cache = self._m3_fc_norm_w + if cache is None or cache[0] is not dm.fc_norm[0].weight: + if torch.cuda.is_current_stream_capturing(): + dispatch.capture_missed("spec_tail", None, "the stacked Eagle3 fc_norm weights") + return None + cache = ( + dm.fc_norm[0].weight, + torch.stack([n.weight.detach() for n in dm.fc_norm]).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], dm.fc_norm[0].variance_epsilon) + hidden = dm.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} + + @staticmethod + def _m3_fc_norm_ok(draft_model: Optional[nn.Module]) -> bool: + """An Eagle3 drafter with an fc over several captured layers, each normalized by its own fc_norm (one + shared epsilon): what the fused first drafter inputs compute.""" + dm = getattr(draft_model, "model", None) + fc_norm = getattr(dm, "fc_norm", None) + return ( + fc_norm is not None + and getattr(dm, "fc", None) is not None + and len({n.variance_epsilon for n in fc_norm}) == 1 + ) + + def _m3_fc_fork( + self, + spec_metadata: Any, + num_tokens: int, + draft_model: Optional[nn.Module], + ) -> Optional[torch.cuda.Stream]: + """Launch the first draft pass's fc_norm + fc on a forked stream beside the target head GEMM (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 or not self._m3_fc_norm_ok(draft_model): + return None + dm = draft_model.model + layers = len(spec_metadata.layers_to_capture) + if ( + dm.hidden_size_in != spec_metadata.hidden_size + or layers != len(dm.fc_norm) + 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() + 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], dm.fc_norm[0].variance_epsilon, use_pdl=False + ) + self._m3_fc_hidden = dm.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 or + torch.compile, 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: Any, + spec_metadata: Any, + lm_head: nn.Module, + draft_model: Optional[nn.Module], + ) -> 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 + ``dispatch.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 + # 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) + <= dispatch.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 self._m3_fc_norm_ok(draft_model) + and fc_norm_w[0] is draft_model.model.fc_norm[0].weight + and len(spec_metadata.layers_to_capture) == fc_norm_w[1].shape[0] + ): + self._m3_capture_norm = (fc_norm_w[1], draft_model.model.fc_norm[0].variance_epsilon) + + @staticmethod + def _m3_sampled_verify_ok(spec_metadata: Any) -> 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: torch.Tensor, + spec_metadata: Any, + batch_size: int, + draft_len: int, + draft_model: Optional[nn.Module] = None, + ): + """Fused verification (see tail_eligible); same outputs as the strict acceptance path.""" + 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_kernel[(1,)]( + 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=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 + ) + 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 (none at FC_EARLY_MAX_ROWS = 0), + # else one beside the head. + fork = self._m3_fc_early_stream + self._m3_fc_early_stream = None + if fork is None: + fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model) + # This rank's vocabulary shard of the target logits: the in-tree LM head's own local BF16 GEMM, whose + # all-gather the fused kernels replace (they exchange per-row records across ranks, not logits). + logits = lm_head.apply_linear(rows, None) + vocab_offset = self.mapping.tp_rank * lm_head.weight.shape[0] + if spec_metadata.is_all_greedy_sample: + result = m3_shard_accept.shard_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + force_total, + self._m3_next_new, + force_vec=force_vec, + ) + else: + n_rows = rows.shape[0] + seeds, offsets = self._rng_state_per_token(spec_metadata, n_rows) + result = m3_sampled_accept.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, + ) + if fork is not None: + self._m3_fc_join(fork) + if result is None: + 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: torch.device + ) -> Optional[torch.Tensor]: + """Per-request forced accepted counts for a fractional forced acceptance, else None. + + Draws exactly as _apply_force_accepted_tokens does for a decode-only batch (same fixed-seed pool, one counter + bump per step, the same (counter, slot) hash), so every TP rank and the generic path agree on the counts. + """ + 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) + # One launch: bump the shared counter and draw, with _apply_force_accepted_tokens' exact arithmetic. + _m3_force_vec_kernel[(1,)]( + 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=max(16, 1 << (batch_size - 1).bit_length()), + ) + return out + + @staticmethod + def _m3_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 + ) + + # ---------------------------------------------------------------------------------------------- fused draft loop + + def prebuild_draft_head(self, draft_model: Optional[nn.Module]) -> bool: + """Build the fragment-ordered copy of the drafter's BF16 head shard that the fused draft loop streams (rebuilt + from the current weight values on every call). Called by the target's ``cache_derived_state`` once the weights + are real, before the executor's warmup, 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.prepare_head(draft_model.lm_head, self.mapping.tp_rank) + return True + + def prebuild_drafter_linears(self, draft_model: Optional[nn.Module]) -> int: + """Put the drafter's BF16 linears on the fixed-tactic low-M kernels (``dispatch.DRAFT_LOWM``) and compile every + covered call, and cap the drafter attention's trtllm-gen KV split per request count + (``dispatch.DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH``). Called by the target's ``cache_derived_state`` once the + weights are real, before the executor's warmup, on every rank. Returns the number of low-M calls compiled.""" + if draft_model is None: + return 0 + dm = getattr(draft_model, "model", None) + attn = getattr(getattr(getattr(dm, "midlayer", None), "self_attn", None), "attn", None) + if attn is not None: + attn.dense_decode_max_seq_len_by_batch = dispatch.DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH + m3_draft_lowm.install(draft_model) + return m3_draft_lowm.prebuild(draft_model) + + def _m3_drafter_supported(self, draft_model: nn.Module) -> bool: + """The drafter the fused draft loop runs: one plain (non-MLA) Eagle3 layer with a gated MLP over the + target-vocabulary embedding (whole table on every rank), a draft vocabulary equal to the target's (no d2t), + all-reduced TP projections, and a bias-free vocab-sharded BF16 head without padding.""" + dm = getattr(draft_model, "model", None) + if ( + dm is None + or getattr(dm, "num_layers", 0) != 1 + or dm._use_mla + or dm._eh_proj_before_attn + ): + return False + layer = dm.midlayer + if layer._next_layer_regular or not isinstance(layer.mlp, GatedMLP): + return False + embed = dm.embed_tokens + if embed is None or getattr(embed, "tp_size", 1) != 1: + return False + if self.mapping is None or self.mapping.tp_size <= 1 or self.mapping.enable_attention_dp: + return False + if getattr(layer.self_attn.o_proj, "all_reduce", None) is None: + return False + if getattr(layer.mlp.down_proj, "all_reduce", None) is None: + return False + if self._d2t is not None or getattr(dm, "d2t", None) is not None: + return False + lm_head = draft_model.lm_head + return self._m3_sharded_head(lm_head) and m3_draft_step.head_supported(lm_head) + + def _m3_fused_loop_ok( + self, + attn_metadata: Any, + spec_metadata: Any, + batch_size: int, + draft_model: Optional[nn.Module], + ) -> 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) + ``dispatch.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 = dispatch.DRAFT_STEP.path + if not self._m3_drafter_supported(draft_model) or not hasattr( + attn_metadata, "kv_lens_cuda" + ): + dispatch.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 dispatch.DRAFT_STEP.covers((draft_len, batch_size)): + dispatch.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: + dispatch.missed( + path, + batch_size, + "the drafter head's fragment-ordered copy is not built for its current weight (prebuild_draft_head)", + ) + return False + return m3_draft_step.workspace(self.mapping, head.wf.device) is not None + + 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: dict[str, Any], + attn_metadata: Any, + spec_metadata: Any, + draft_model: nn.Module, + draft_kv_cache_manager: Any, + batch_size: int, + num_accepted_tokens: torch.Tensor, + ) -> torch.Tensor: + """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_layernorm 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 in-tree attention and MLP, each followed by the in-tree all-reduce with the residual add + and the next RMSNorm fused into it (post_attention_layernorm, 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 + dm = draft_model.model + layer = dm.midlayer + attn = layer.self_attn + mlp = layer.mlp + attn_all_reduce = attn.o_proj.all_reduce + mlp_all_reduce = mlp.down_proj.all_reduce + table = dm.embed_tokens.weight + 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:] + + no_all_reduce = AllReduceParams(enable_allreduce=False) + 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, + layer.input_layernorm.weight, + layer.input_layernorm.variance_epsilon, + layer.hidden_norm.weight, + layer.hidden_norm.variance_epsilon, + ) + h = attn( + position_ids=pos, + hidden_states=x, + attn_metadata=md, + all_reduce_params=no_all_reduce, + ) + h, residual = draft_ar_norm( + attn_all_reduce, h, residual, layer.post_attention_layernorm + ) + h = mlp(h, final_all_reduce_params=no_all_reduce) + h, to_save = draft_ar_norm(mlp_all_reduce, h, residual, dm.norm) + if dm._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.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 + + +class M3SpecDecForwardMixin: + """The target CausalLM's step forward with the fused-verify hooks of the speculative wrapper. + + Mirrors ``SpecDecOneEngineForCausalLM.forward`` (tensorrt_llm/_torch/models/modeling_speculative.py) with the + fused-verify branches: before the target forward the worker arms the capture norm, and a step the worker verifies + from the logits shard (``M3Eagle3Worker.tail_eligible``) skips the gathered target logits. + + The class it is mixed into (listed before ``DecoderModelForCausalLM``) provides ``model`` (called like the in-tree + decoder model, capturing through ``M3Eagle3Worker.capture_hidden_states``), ``lm_head``, ``logits_processor``, + ``model_config``, ``draft_model`` and ``spec_worker`` (an ``M3Eagle3Worker``, or None without speculation). + """ + + def forward( + self, + attn_metadata: Any, + 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[Any] = None, + resource_manager: Optional[Any] = None, + **kwargs, + ) -> Any: + spec_worker: Optional[M3Eagle3Worker] = self.spec_worker + if spec_worker is not None and spec_metadata is not None: + spec_worker.arm_capture_norm( + attn_metadata, spec_metadata, self.lm_head, self.draft_model + ) + hidden_states = self.model( + input_ids=input_ids, + attn_metadata=attn_metadata, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + **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): + if spec_worker is not None: + spec_worker.capture_hidden_states( + spec_metadata, FINAL_HIDDEN_LAYER_ID, hidden_states, None + ) + else: + capture_hidden_states( + spec_metadata, FINAL_HIDDEN_LAYER_ID, hidden_states, None, None + ) + if attn_metadata.padded_num_tokens is not None: + hidden_states = hidden_states[: attn_metadata.num_tokens] + + if spec_worker is not None: + # The fused verify computes this rank's logits shard itself and needs no gathered logits. + if spec_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 + ) + return spec_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, + ) + + return self.logits_processor.forward( + hidden_states, + self.lm_head, + attn_metadata, + return_context_logits, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/weights.py new file mode 100644 index 000000000000..c69fa3be8ef1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/weights.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight loading: MiniMax-M3 NVFP4 / sm_103 / tp8. + +The checkpoint is the vision-language wrapper's. ``language_model.*`` holds the +language model; ``vision_tower.*``, ``multi_modal_projector.*`` and +``patch_merge_mlp.*`` hold the vision tower and its projector. This target is +text only, so the keys split three ways: + +* ``language_model.*`` goes to the built-in MiniMax-M3 text loader with the + prefix stripped (its multi-token-prediction head, ``model.mtp.*``, is not + loaded, as in the built-in model). That loader packs the five-way + Q/K/V/index-Q/index-K projection, folds the Gemma boundary norms and shards + every tensor for tp8; the modules it fills are this target's subclasses of + the built-in ones, under the built-in names. +* The vision tower and the projector are a predicted non-load: listed here, + never read. +* Any other key fails the load, naming it, rather than being dropped. +""" + +from typing import Any, Dict, Optional + +from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper +from tensorrt_llm._torch.models.modeling_minimaxm3 import ( + MiniMaxM3ForCausalLM, + _strip_language_model_prefix, +) + +_LANG_PREFIX = "language_model." + +# The vision tower's and the projector's key families: in the checkpoint, not loaded by this text-only target. +PREDICTED_NON_LOAD = ("vision_tower.", "multi_modal_projector.", "patch_merge_mlp.") + + +def load( + model, + weights: Dict[str, Any], + weight_mapper: Optional[BaseWeightMapper] = None, + params_map: Optional[Dict[str, str]] = None, + allow_partial_loading: bool = False, +) -> None: + """Load the language model's weights into ``model``, the target shell, and check every other key is predicted.""" + unknown = sorted( + k + for k in weights.keys() + if not k.startswith(_LANG_PREFIX) and not k.startswith(PREDICTED_NON_LOAD) + ) + assert not unknown, ( + f"{len(unknown)} checkpoint key(s) are neither language-model weights nor a predicted non-load, " + f"first {unknown[:5]}" + ) + text_weights, _ = _strip_language_model_prefix( + {k: v for k, v in weights.items() if k.startswith(_LANG_PREFIX)} + ) + assert text_weights, f"the checkpoint has no {_LANG_PREFIX}* keys" + MiniMaxM3ForCausalLM.load_weights( + model, + text_weights, + weight_mapper=weight_mapper, + params_map=params_map, + allow_partial_loading=allow_partial_loading, + ) 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 480ea23f9b02..c6b166dc1300 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 @@ -67,20 +67,21 @@ 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), - so it selects a target rather than a runtime branch. + 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. """ - if not ( - m.world_size == 4 - and m.tp_size == 4 - and m.pp_size == 1 - and m.cp_size == 1 - and m.moe_tp_size == 4 - and m.moe_ep_size == 1 - and not m.enable_attention_dp - ): - return None - return "tp4" + for tp in (4, 8): + if ( + m.world_size == tp + and m.tp_size == tp + and m.pp_size == 1 + and m.cp_size == 1 + and m.moe_tp_size == tp + and m.moe_ep_size == 1 + and not m.enable_attention_dp + ): + return f"tp{tp}" + return None def _eagle3_one_model_or_none(spec_config: Optional[Any]) -> bool: @@ -137,13 +138,26 @@ def route( return ModelingV2MiniMaxM3Nvfp4Sm103Tp4 + # The same checkpoint on 8 GB300s over two trays (sm103), tensor-parallel + # attention and experts, one-model Eagle3 or no speculative decoding. + if ( + compute_capability == _SM + and shape == _SHAPE + and quant_algo == QuantAlgo.MIXED_PRECISION + and parallel == "tp8" + and speculation + ): + from .minimax_m3_nvfp4__sm_103__tp8.modeling import ModelingV2MiniMaxM3Nvfp4Sm103Tp8 + + return ModelingV2MiniMaxM3Nvfp4Sm103Tp8 + raise NoTarget( f"compute capability {compute_capability}, text_config shape {shape}, " f"quantization {quant_algo.name if quant_algo is not None else 'none'}, parallel " f"ws={mapping.world_size} tp={mapping.tp_size} pp={mapping.pp_size} cp={mapping.cp_size} " f"moe_tp={mapping.moe_tp_size} moe_ep={mapping.moe_ep_size} " f"attention_dp={mapping.enable_attention_dp}, speculative decoding " - f"{spec_config.spec_dec_mode.name if spec_config is not None else 'off'}; the one " - "minimax_m3 target is the MiniMax-M3 NVFP4 checkpoint on sm103 (GB300), 4 ranks of " + f"{spec_config.spec_dec_mode.name if spec_config is not None else 'off'}; 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" ) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py index 1f43ae3bb385..b7e0a50c60a3 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py @@ -62,6 +62,7 @@ from tensorrt_llm.quantization.mode import QuantMode from .interface import FmhaPhase +from .msa_decode import dense_decode_max_seq_len from .phased import FmhaParams, PhasedFmha from .utils import ( get_attention_chunk_size, @@ -1163,6 +1164,8 @@ def run_generation( bmm1_scale if fp8_context_fmha and bmm1_scale is not None else bmm1_scale_static ) gen_bmm2_scale = bmm2_scale if fp8_context_fmha and bmm2_scale is not None else 1.0 + # A layer may cap the KV split's max_seq_len per request count (a no-op unless it sets the cap). + max_kv_len = dense_decode_max_seq_len(attn, max_kv_len, batch_beam) flashinfer.decode.trtllm_batch_decode_with_kv_cache( query=q_processed, diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/_m3_nvfp4_cache.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/attention/test_modeling_v2_m3_dense_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_dense_decode.py new file mode 100644 index 000000000000..d7d912e26bfd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_dense_decode.py @@ -0,0 +1,152 @@ +# 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_103__tp8 import ( # noqa: E501 + dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.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_103__tp8.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) + + +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. The Eagle3 drafter layer's + decode (the same 8 query heads over one KV head per rank at TP8) takes the same cap, so the same cells hold it.""" + assert dict(dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH) == {1: CAPPED} + assert dict(dispatch.DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH) == {1: CAPPED} + assert (HEADS, KV_HEADS) == (8, 1) + 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_103__tp8/attention/test_modeling_v2_m3_fused_score_select.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_fused_score_select.py new file mode 100644 index 000000000000..9b1edb6f77dd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_fused_score_select.py @@ -0,0 +1,230 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_fused_score_select entry: MSA index scoring and top-16 block selection of a decode +span of at most two requests in one launch (``torch.ops.trtllm.minimax_m3_tp8_fused_score_select``). + +Every cell is held bit for bit to the entry's ``reference`` (``m3_select_long.reference`` over the per-KV-head maxima +of ``m3_index_score.reference``) and to the pipeline it replaces in the target: this target's ``m3_index_score`` (the +prefetching kernel at ``SPLIT_K``), the per-KV-head amax and ``m3_select_long``, on the same inputs. The selection is a +table of block ids, so the claim is exact. +""" + +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_103__tp8 import ( # noqa: E501 + dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 + m3_fused_score_select as fss, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_fused_score_select import ( # noqa: E501 + MAX_REQUESTS, + scratch_shape, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_fused_score_select import ( # noqa: E501 + m3_fused_score_select as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_index_score import ( # noqa: E501 + SPLIT_K, + m3_index_score, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, + m3_select_long, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fused_score_select requires a CUDA device" + +PAGE = 128 +HEADS = 1 # index heads per rank at TP8 +HEAD_DIM = 128 + + +def _n_valid(seq_lens, qlen: int) -> torch.Tensor: + """Each query row's causal valid-block count, as the metadata stages it (``per_token_valid_blocks``).""" + vals = [ + (max(seq - qlen + offset + 1, 0) + PAGE - 1) // PAGE + for seq in seq_lens + for offset in range(qlen) + ] + return torch.tensor(vals, dtype=torch.int32, device="cuda") + + +def _inputs(spec, seed: int, dtype=torch.float8_e4m3fn): + """Random FP8 index Q [T, 1, 128] and paged index-K [pages, 128, 128] behind a shuffled block table (as the index + scorer's test builds them), the request lengths and the rows' valid-block counts.""" + seq_lens, qlen = spec["seq_lens"], spec["qlen"] + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + max_blocks = (max(seq_lens) + PAGE - 1) // PAGE + 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, _n_valid(seq_lens, qlen) + + +def _out(total_q: int, num_kv_heads: int) -> torch.Tensor: + """The [total_q, num_kv_heads, 16] view of a head-major backing, as the target hands it over.""" + return torch.full((num_kv_heads, total_q, TOPK), -7, dtype=torch.int32, device="cuda").permute( + 1, 0, 2 + ) + + +def _args(spec, inputs, out, scratch, counters): + idx_q, index_k, block_table, seq, n_valid = inputs + return ( + idx_q, + index_k, + block_table, + seq, + n_valid, + out, + scratch, + counters, + spec["qlen"], + spec["num_kv_heads"], + spec["init_blocks"], + spec["local_blocks"], + spec["max_tiles"], + False, + ) + + +def _workspace(max_tiles: int): + """Scratch filled with garbage (the op needs no initialization) and zeroed counters.""" + scratch = torch.full(scratch_shape(max_tiles), float("nan"), device="cuda") + counters = torch.zeros(MAX_REQUESTS, dtype=torch.int32, device="cuda") + return scratch, counters + + +def _pipeline(spec, inputs) -> torch.Tensor: + """What the target runs on a wider span: this target's index scorer (prefetching kernel), the per-KV-head amax over + its index heads, and the long-row selector.""" + idx_q, index_k, block_table, seq, n_valid = inputs + total_q = idx_q.shape[0] + max_tiles, kv = spec["max_tiles"], spec["num_kv_heads"] + backing = torch.full( + (HEADS, max_tiles, total_q), -float("inf"), dtype=torch.float32, device="cuda" + ) + m3_index_score( + idx_q, index_k, block_table, seq, backing.transpose(1, 2), spec["qlen"], SPLIT_K, False + ) + per_kv = backing.view(kv, HEADS // kv, max_tiles, total_q).amax(dim=1) + table = _out(total_q, kv) + return m3_select_long( + per_kv, n_valid, table, spec["init_blocks"], spec["local_blocks"], False, 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 + spec = cell.spec + inputs = _inputs(spec, zlib.crc32(cell.why.encode()) % 2**31) + assert len(spec["seq_lens"]) == spec["batch"] + total_q = inputs[0].shape[0] + out = _out(total_q, spec["num_kv_heads"]) + scratch, counters = _workspace(spec["max_tiles"]) + args = _args(spec, inputs, out, scratch, counters) + certifying(op) + op(*args) + ref = op.reference(*args) + torch.cuda.synchronize() + + assert out.dtype is torch.int32 and out.shape == (total_q, spec["num_kv_heads"], TOPK) + assert torch.equal(out, ref), "the selected blocks differ from the reference" + assert not bool(counters.any()), "the arrival counters must be back at zero after the call" + assert torch.equal(out, _pipeline(spec, inputs)), ( + "the selected blocks differ from m3_index_score + per-KV amax + m3_select_long" + ) + + # A second call on the same scratch and counters (as the next sparse layer makes) selects the same. + again = _out(total_q, spec["num_kv_heads"]) + op(*_args(spec, inputs, again, scratch, counters)) + torch.cuda.synchronize() + assert torch.equal(again, out) and not bool(counters.any()) + + +@pytest.mark.parametrize("cell", op.CELLS[:2], ids=[c.why[:44] for c in op.CELLS[:2]]) +def test_the_target_wrapper_is_the_entry(cell) -> None: + """``kernels/m3_fused_score_select``: the geometry check admits the cell, the scratch and counters are created + once per (device, engine block count), and the wrapper's table is the entry's.""" + spec = cell.spec + inputs = _inputs(spec, 3) + idx_q, index_k, block_table, seq, n_valid = inputs + max_tiles = spec["max_tiles"] + assert fss.supported(idx_q, index_k, spec["qlen"], spec["num_kv_heads"], max_tiles) is None + bufs = fss.buffers(idx_q.device, max_tiles) + assert bufs is not None and fss.buffers(idx_q.device, max_tiles) is bufs + table = fss.fused_score_select( + idx_q, + index_k, + block_table, + seq, + n_valid, + decode_query_len=spec["qlen"], + num_kv_heads=spec["num_kv_heads"], + init_blocks=spec["init_blocks"], + local_blocks=spec["local_blocks"], + max_tiles=max_tiles, + scratch=bufs[0], + counters=bufs[1], + ) + torch.cuda.synchronize() + ref = op.reference(*_args(spec, inputs, _out(idx_q.shape[0], 1), *_workspace(max_tiles))) + assert torch.equal(table, ref) + assert not bool(bufs[1].any()) + + +def test_the_geometry_check_refuses_what_the_kernel_does_not_take() -> None: + spec = op.CELLS[0].spec + idx_q, index_k, *_ = _inputs(spec, 5) + assert "not FP8" in fss.supported(idx_q.bfloat16(), index_k, 4, 1, 8192) + assert "geometry" in fss.supported(idx_q, index_k, 5, 1, 8192) + wide_q = idx_q.expand(-1, 4, -1).contiguous() + assert "geometry" in fss.supported(wide_q, index_k, 4, 1, 8192) + assert "stage ring" in fss.supported(idx_q, index_k, 4, 1, 16384) + + +def test_the_dispatch_range_is_the_kernel_limit() -> None: + assert dispatch.FUSED_SCORE_SELECT_MAX_REQUESTS == MAX_REQUESTS + assert {c.spec["batch"] for c in op.CELLS} == set(dispatch.FUSED_SCORE_SELECT.keys) + + +def test_guard_refuses_nonzero_counters() -> None: + spec = op.CELLS[0].spec + inputs = _inputs(spec, 7) + scratch, counters = _workspace(spec["max_tiles"]) + counters[0] = 1 + with pytest.raises(AssertionError, match="counters must be zero"): + with validating(op): + op(*_args(spec, inputs, _out(inputs[0].shape[0], 1), scratch, counters)) + + +def test_guard_refuses_a_request_past_its_block_table() -> None: + spec = op.CELLS[4].spec + idx_q, index_k, block_table, seq, n_valid = _inputs(spec, 9) + seq = seq.clone() + seq[1] = block_table.shape[1] * PAGE + 1 + scratch, counters = _workspace(spec["max_tiles"]) + with pytest.raises(AssertionError, match="past its block-table row"): + with validating(op): + op( + *_args( + spec, + (idx_q, index_k, block_table, seq, n_valid), + _out(idx_q.shape[0], 1), + scratch, + counters, + ) + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_index_score.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_index_score.py new file mode 100644 index 000000000000..090e82cb7020 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_index_score.py @@ -0,0 +1,94 @@ +# 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 prefetch. + +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 must not +change. The tp8 target runs the prefetching kernel only (spans of more than two requests; ``early`` is never +selected), so that is the schedule driven here. +""" + +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_103__tp8.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_103__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 +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" + + +def test_bf16_index_is_the_in_tree_scorer() -> None: + idx_q, index_k, block_table, seq, stride = _inputs( + (5000, 300, 2049), 4, 11, dtype=torch.bfloat16 + ) + score, stock = _score(stride, 12), _score(stride, 12) + op(idx_q, index_k, block_table, seq, score, 4, SPLIT_K, False) + 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, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_producer.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_producer.py new file mode 100644 index 000000000000..a948a445f648 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_producer.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 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_103__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_103__tp8.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, 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)) + + +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_103__tp8/attention/test_modeling_v2_m3_select_long.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_select_long.py new file mode 100644 index 000000000000..96ad7bb7b05c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_select_long.py @@ -0,0 +1,159 @@ +# 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_103__tp8.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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 == "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) + + +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_103__tp8/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_sparse_decode.py new file mode 100644 index 000000000000..3551a46ca09c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_sparse_decode.py @@ -0,0 +1,144 @@ +# 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 + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_sparse_decode import ( # noqa: E501 + m3_sparse_decode as op, +) + +__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 + +assert torch.cuda.is_available(), "m3_sparse_decode requires a CUDA device" + +# Band, in BF16 ulps (2^-8) of each [8, 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. None of these depends on the head count: +# measured on GB300 over the same cells at the tp4 target's 16 heads, 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 test_a_selected_step_the_lc8_kernel_declines_is_a_miss() -> None: + """A selected table at the TP8 heads is the one-launch kernel's (dispatch.SPARSE_DECODE_LC8): an output it cannot + write (here a strided view) raises under require instead of silently taking the split-K chain.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( + dispatch, + ) + + case = build_step((9000, 40000), 4, identity=False, seed=31) + wide = torch.empty( + case.q.shape[0], NUM_HEADS, 2 * HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + output = wide[..., :HEAD_DIM] + assert not output.is_contiguous() + dispatch.reset_logged() + dispatch.set_require(True) + try: + with pytest.raises(dispatch.FastPathUnavailable, match="m3_sparse_decode_lc8"): + _call(case, output, False) + finally: + dispatch.set_require(False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py new file mode 100644 index 000000000000..39766f50d67f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py @@ -0,0 +1,195 @@ +# 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_103__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_103__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 + +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)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_msa_kv_update.py new file mode 100644 index 000000000000..ffca6d60eefd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp8.kernels.msa_kv_update import ( # noqa: E501 + msa_kv_update as hook, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__tp8.ops.attention.msa_kv_update import ( # noqa: E501 + msa_kv_update as op, +) + +__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_103__tp8/comm/_m3_draft_ar_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_draft_ar_norm_op_matrix.py new file mode 100644 index 000000000000..271bc5d6fd3a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_draft_ar_norm_op_matrix.py @@ -0,0 +1,394 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fast draft loop's layer boundaries (``spec_worker.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 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_draft_ar_norm.py``. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_draft_ar_norm_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 _target(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """The draft loop's boundary: ``spec_worker.draft_ar_norm`` with the drafter's norm module.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.spec_worker import ( # noqa: E501 + draft_ar_norm, + ) + from tensorrt_llm._torch.modules.rms_norm import RMSNorm + + norm = RMSNorm(hidden_size=HIDDEN, eps=EPS, dtype=torch.bfloat16).cuda() + with torch.no_grad(): + norm.weight.copy_(w) + return draft_ar_norm(ar, x, residual.clone(), norm) + + +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) + 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 + + 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_103__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..d322305df03b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,670 @@ +# 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 GB300 trays, which one 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``). + +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 + +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) +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_103__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): + 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) + _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 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_103__tp8.kernels.m3_mc_allreduce as _ar # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__tp8.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 + + 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_103__tp8/comm/_rank_job.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_rank_job.py new file mode 100644 index 000000000000..72e49f20a1e6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_rank_job.py @@ -0,0 +1,107 @@ +# 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, and a source checkout is not necessarily installed. + # tests/unittest/_torch/modeling_v2/models/minimax_m3//comm -> repo root. + repo_root = Path(__file__).resolve().parents[8] + assert (repo_root / "tensorrt_llm").is_dir(), ( + f"expected the repo root at {repo_root}, found no tensorrt_llm/ there; " + f"this file moved without its parents[] index following" + ) + env["PYTHONPATH"] = os.pathsep.join(p for p in (str(repo_root), env.get("PYTHONPATH", "")) if p) + + 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_103__tp8/comm/test_modeling_v2_m3_draft_ar_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_draft_ar_norm.py new file mode 100644 index 000000000000..d0d521428d12 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_draft_ar_norm.py @@ -0,0 +1,38 @@ +# 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 launcher; see ``_rank_job``. It requires +``spec_worker.draft_ar_norm`` to equal stock Eagle3's boundary (plain all-reduce, then the ``RMSNorm`` module with the +residual) 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 + +assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" + +_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_draft_ar_norm_is_stock_eagle3(world_size) -> None: + out = _rank_job.run("m3_draft_ar_norm", 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_103__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_103__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..31ebc26c4a6d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,60 @@ +# 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; 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 GB300 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_103__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_103__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, + ): + 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_103__tp8/gemm/test_modeling_v2_m3_draft_lowm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_draft_lowm.py new file mode 100644 index 000000000000..a8eb5a5c1425 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_draft_lowm.py @@ -0,0 +1,196 @@ +# 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. + +Every cell is held to the fp64 reference within the derived band. The target's linear method +(``kernels/m3_draft_lowm.DraftLowMLinearMethod``) is held on a real ``Linear`` (BF16, no bias): a covered row count is +the entry's call bit for bit, an uncovered one the in-tree method's. +""" + +import zlib + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 + dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 + m3_draft_lowm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm.m3_draft_lowm import ( # noqa: E501 + TACTICS, + tactic_for, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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 +from tensorrt_llm._torch.modules.linear import Linear, UnquantizedLinearMethod + +__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 + +DRAFTER_LINEARS = { + (1280, 12288): "qkv", + (6144, 1024): "o_proj", + (4608, 6144): "gate_up", + (6144, 2304): "down", + (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_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 _linear(nk, seed: int, bias: bool = False) -> Linear: + n, k = nk + linear = Linear(k, n, bias=bias, dtype=torch.bfloat16).cuda() + _, w = build(n, k, 1, seed) + linear.weight.data.copy_(w) + if bias: + linear.bias.data.zero_() + return linear + + +class _Drafter(nn.Module): + """A module tree holding drafter linears, as ``install`` / ``prebuild`` walk it.""" + + def __init__(self, *linears: Linear): + super().__init__() + self.linears = nn.ModuleList(linears) + + +@pytest.mark.parametrize("nk", sorted(DRAFTER_LINEARS), ids=lambda nk: DRAFTER_LINEARS[nk]) +def test_the_linear_method_is_the_entry_and_falls_back_past_its_rows(nk) -> None: + n, k = nk + linear = _linear(nk, seed=n + k) + drafter = _Drafter(linear) + assert m3_draft_lowm.install(drafter) == 1 + assert isinstance(linear.quant_method, m3_draft_lowm.DraftLowMLinearMethod) + m3_draft_lowm.prebuild(drafter) + use_pdl = get_env_enable_pdl() + last = TACTICS[nk][-1][0] + dispatch.set_require(True) + for m in range(1, last + 1): + assert dispatch.DRAFT_LOWM.covers((nk, m)) + x, _ = build(n, k, m, seed=m) + got = linear.quant_method.apply(linear, x, None) + kind, tactic = tactic_for(n, k, m) + want = op( + x, linear.weight, 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 + assert torch.equal(linear(x).view(torch.int16), want.view(torch.int16)), m + # Past the table: the in-tree method, silently, even under require. + for m in (last + 1, 2 * last + 3): + assert not dispatch.DRAFT_LOWM.covers((nk, m)) + x, _ = build(n, k, m, seed=m) + got = linear.quant_method.apply(linear, x, None) + stock = UnquantizedLinearMethod().apply(linear, x, None) + assert torch.equal(got.view(torch.int16), stock.view(torch.int16)), m + + +def test_install_takes_only_covered_bias_free_bf16_linears() -> None: + covered = _linear((6144, 1024), seed=1) + biased = _linear((6144, 1024), seed=2, bias=True) + other = _linear((256, 512), seed=3) + drafter = _Drafter(covered, biased, other) + assert m3_draft_lowm.install(drafter) == 1 + assert isinstance(covered.quant_method, m3_draft_lowm.DraftLowMLinearMethod) + assert type(biased.quant_method) is UnquantizedLinearMethod + assert type(other.quant_method) is UnquantizedLinearMethod + # A second install finds nothing left to take. + assert m3_draft_lowm.install(drafter) == 0 + + +def test_a_covered_call_not_prebuilt_raises_under_require(monkeypatch) -> None: + monkeypatch.setattr(m3_draft_lowm, "_BUILT", set()) + linear = _linear((6144, 1024), seed=4) + m3_draft_lowm.install(_Drafter(linear)) + x, _ = build(6144, 1024, 1, seed=5) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="not compiled before warmup"): + linear.quant_method.apply(linear, x, None) + # Outside require it falls back to the in-tree method. + dispatch.set_require(False) + got = linear.quant_method.apply(linear, x, None) + stock = UnquantizedLinearMethod().apply(linear, x, None) + assert torch.equal(got.view(torch.int16), stock.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_103__tp8/gemm/test_modeling_v2_m3_mx_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mx_gemv.py new file mode 100644 index 000000000000..960f824f877f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mx_gemv.py @@ -0,0 +1,91 @@ +# 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_103__tp8.ops.gemm.m3_mx_gemv import ( # noqa: E501 + O_PROJ, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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 tp4 target's o_proj and QKV (6 seeds each, GB300; K 2048 and 6144, the tp8 projections' K is 1024..6144), 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.""" + n, k = O_PROJ + x, _, _, w, wsf = build(n, k, 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, k), xsf, w, wsf) + + +def test_guard_refuses_a_strided_weight() -> None: + n, k = O_PROJ + _, xq, xsf, w, wsf = build(n, k, 1, seed=3) + wide = torch.zeros(n, 2 * k, dtype=w.dtype, device="cuda") + wide[:, :k] = w + with pytest.raises(AssertionError, match="contiguous e4m3"): + with validating(op): + op(xq, xsf, wide[:, :k], wsf) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..c88f84729b45 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py @@ -0,0 +1,67 @@ +# 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_103__tp8.ops.gemm.m3_mx_gemv import ( # noqa: E501 + O_PROJ, + m3_mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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 (a served shape: + the same x, contiguous, is computed).""" + n, k = O_PROJ + x, w, wsf = build(n, k, 2, seed=1) + wide = torch.zeros(2, 2 * k, dtype=torch.bfloat16, device="cuda") + wide[:, :k] = x + assert op(wide[:, :k], w, wsf) is None + assert op(wide[:, :k].contiguous(), w, wsf) is not None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/_m3_moe_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/_m3_moe_operands.py new file mode 100644 index 000000000000..80ef9fbd7225 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/moe/test_modeling_v2_m3_moe_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_moe_engine.py new file mode 100644 index 000000000000..87655f1fc9ac --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_moe_engine.py @@ -0,0 +1,134 @@ +# 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 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_103__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_103__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" + +_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_103__tp8/moe/test_modeling_v2_m3_router_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_gemv.py new file mode 100644 index 000000000000..e85564fc0fb4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_gemv.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the router entries of the MiniMax-M3 tp8 target: m3_router_gemv, m3_router_gemv_fp4 and +m3_router_tc (one kernel source, one test module).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__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_103__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_103__tp8.ops.moe.m3_router_tc import ( # noqa: E501 + m3_router_tc as tc, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "the router entries require a CUDA device" + +#: 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" + + +@pytest.mark.parametrize("cell", gemv.CELLS, ids=[c.why[:48] for c in gemv.CELLS]) +def test_gemv_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], spec["e"], spec["k"], _seed(cell)) + with validating(gemv): + out = gemv(x, w, 1, False) + _compare_logits(out, gemv.reference(x, w, 1, False)) + + +@pytest.mark.parametrize("cell", gemv_fp4.CELLS, ids=[c.why[:48] for c in gemv_fp4.CELLS]) +def test_gemv_fp4_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], spec["e"], spec["k"], _seed(cell)) + sf_scale = torch.tensor([spec["sf_scale"]], dtype=torch.float32, device="cuda") + with validating(gemv_fp4): + logits, q, sf = gemv_fp4(x, w, sf_scale, False) + ref_logits, _, _ = gemv_fp4.reference(x, w, sf_scale, False) + _compare_logits(logits, ref_logits) + _compare_fp4(q, sf, x, sf_scale) + + +@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_the_gemv_is_a_true_fp32_product() -> None: + """And the GEMV, by contrast, is no TF32 product: the TF32 reference rejects it.""" + x, w = _build(1, 128, 6144, 6) + out = gemv(x, w, 1, False) + with pytest.raises(AssertionError): + _compare_logits(out, tc.reference(x, w, None, False)[0]) + + +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) + + +def test_guard_refuses_a_zero_global_scale() -> None: + x, w = _build(1, 128, 6144, 4) + with pytest.raises(AssertionError, match="positive and finite"): + with validating(gemv_fp4): + gemv_fp4(x, w, torch.zeros(1, dtype=torch.float32, device="cuda"), False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py new file mode 100644 index 000000000000..d7a3f6fef547 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py @@ -0,0 +1,71 @@ +# 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_103__tp8.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + VOCAB, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/_m3_spec_inputs.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/_m3_spec_inputs.py new file mode 100644 index 000000000000..42fdec65c754 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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 tp8 target's spec entry tests: TP8 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_103__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 (25008 rows) 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_103__tp8/spec/test_modeling_v2_m3_capture_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_capture_norm.py new file mode 100644 index 000000000000..8720f10a3168 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/spec/test_modeling_v2_m3_chunk_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_chunk_norm.py new file mode 100644 index 000000000000..9a0fad240f3e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_input.py new file mode 100644 index 000000000000..794deaf1678d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp8.kernels import ( # noqa: E501 + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__tp8/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_step.py new file mode 100644 index 000000000000..3e76283f4d90 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_step.py @@ -0,0 +1,611 @@ +# 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 as BF16 fragments, +the cross-rank argmax over the eight TP8 ranks 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. The tp8 target never builds a +coded copy of the head: the op's coded stream (``exc_off`` given) is not certified here. +""" + +import warnings +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 + m3_draft_step as draft_step_kernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.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_103__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_103__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 + + +#: The op's trailing coded-head arguments (exc_off, exc_pos, exc_w, exp_base): none, the BF16 fragments. +NO_CODING = (None, None, None, 0) + + +def _head(w: torch.Tensor): + """(wf, NO_CODING) of a BF16 shard: its BF16 fragments, as ``prepare_head`` builds them.""" + return bf16_fragments(w), NO_CODING + + +def _args( + x, + x_rows, + nreq, + wf, + rank, + world, + step, + g, + uc=None, + keys=None, + counters=None, + num_acc=None, + coded=NO_CODING, +): + """The op's argument tuple for one rank (fresh workspace and bookkeeping buffers unless given); ``coded`` = the + op's trailing coded-head arguments, always ``NO_CODING`` here.""" + 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) + assert not spec["coded"], "the tp8 target certifies the BF16 fragments only" + wf, coded = _head(randn((vs, k), g, 0.02)) + 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) + assert not hasattr(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, step: int = 1) -> list: + """Each rank's fused draft tokens: every rank's own key from a one-rank run, then each rank's eight-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) 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, 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)): + 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 + + +def test_forced_bf16_ties_take_the_lowest_global_id() -> 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, + 7 * DRAFT_SHARD + 1, + ): + w[gid, :4] = win + _assert_stock(w, _ones_rows(5), 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) + + +def test_signed_zero_maxima_take_the_lowest_global_id() -> 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 + # Zeros by cancellation on ranks 1, 2, 3 and 6; 49999 (rank 1's 25008..50015) is the lowest of them. + for gid in (2 * DRAFT_SHARD + 5, 6 * DRAFT_SHARD + 4000, 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, 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"]) +def test_near_equal_top2_follow_the_bf16_logits(lower, higher, winner, where) -> 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 6; 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 + 24000, 6 * 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), 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, 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) + 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"]) +def test_top2_at_a_bf16_midpoint(lower, higher, winner, where) -> 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 6, 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 + 24000, 6 * 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, expect) + + +# ------------------------------------------------------------------ the head copy prepare_head builds + + +def test_prepare_head_streams_the_bf16_fragments_at_every_k() -> None: + """prepare_head keeps the fragment-ordered BF16 copy for every shard, a K the tp4 coding took (a multiple of 64) + included; the op's head weight is the shard bit for bit, and a coded stream is refused.""" + g = gen(seed_of("prepare bf16")) + for k in (1024, 96): + w = randn((4096, k), g, 0.02) + prepared = draft_step_kernels.prepare_head(SimpleNamespace(weight=w, bias=None), 2) + assert prepared.vocab_offset == 2 * 4096 and prepared.weight is w + assert torch.equal(prepared.wf, bf16_fragments(w)) + assert torch.equal(head_weight(prepared.wf, k).view(torch.int16), w.view(torch.int16)) + with pytest.raises(NotImplementedError): + head_weight(prepared.wf, 96, torch.zeros(2, dtype=torch.int32, device="cuda")) + assert not hasattr(draft_step_kernels.dispatch, "DRAFT_HEAD_CODED") + assert all(not c.spec["coded"] for c in op.CELLS) + + +def test_the_cells_are_the_tp8_head() -> None: + assert (WORLD, DRAFT_SHARD, HIDDEN) == (8, 25008, 6144) + for cell in op.CELLS: + assert (cell.spec["world"], cell.spec["vs"], cell.spec["k"]) == (WORLD, DRAFT_SHARD, HIDDEN) + assert 0 <= cell.spec["rank"] < WORLD + + +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_103__tp8/spec/test_modeling_v2_m3_exact_numerics.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_exact_numerics.py new file mode 100644 index 000000000000..3916557e39a5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_exact_numerics.py @@ -0,0 +1,305 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tp8 target's speculative step against the in-tree (stock) Eagle3 numerics, at the real TP8 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 BF16 linears on fixed-tactic BF16 low-M + kernels at the rows ``dispatch.DRAFT_LOWM`` covers, held to the fp64 product and to the entry by + test_modeling_v2_m3_draft_lowm.py); 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 (8 x 25008 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_103__tp8 import ( # noqa: E501 + linears, + spec_worker, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 + m3_draft_lowm, + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.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_103__tp8.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, UnquantizedLinearMethod +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 TP8: an LM head over that rank's + 25008 rows (the TP8 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 eight 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 6 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[6 * 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, the drafter linears' low-M method is the in-tree BF16 method's subclass, and the fused head + streams the BF16 head weight itself (no coded copy).""" + assert ( + spec_worker.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(spec_worker.M3Eagle3Worker, name) is getattr(Eagle3OneModelWorker, name), ( + name + ) + for module, name in ( + (spec_worker, "M3DraftGemvLinearMethod"), + (spec_worker, "adopt_draft_linears"), + (spec_worker, "greedy_head_outputs"), + (linears, "M3FP8LMHeadMethod"), + (linears, "adopt_fp8_lm_heads"), + ): + assert not hasattr(module, name), f"{module.__name__}.{name}" + assert issubclass(m3_draft_lowm.DraftLowMLinearMethod, UnquantizedLinearMethod) + assert not hasattr(spec_worker.M3Eagle3Worker, "prebuild_target_head") + pkg = spec_worker.__name__.rsplit(".", 1)[0] + for gone in ( + "kernels.m3_draft_head", + "ops.spec.m3_draft_head", + "ops.gemm.m3_draft_gemv", + "kernels.m3_target_head", + "ops.gemm.m3_target_head", + ): + 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_103__tp8/spec/test_modeling_v2_m3_sampled_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_sampled_accept.py new file mode 100644 index 000000000000..25a1db2b0943 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_sampled_accept.py @@ -0,0 +1,167 @@ +# 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_103__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_103__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" + + +@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) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_shard_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_shard_accept.py new file mode 100644 index 000000000000..e32448e55119 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__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_103__tp8/spec/test_modeling_v2_m3_slot_scatter.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_slot_scatter.py new file mode 100644 index 000000000000..ba3220639b65 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/spec/test_modeling_v2_m3_step_gather.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_step_gather.py new file mode 100644 index 000000000000..e1a5aed2034b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/test_modeling_v2_m3_capture_boundary.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_capture_boundary.py new file mode 100644 index 000000000000..ce69ccb0a568 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_capture_boundary.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decoder layer's feed-forward exit on Eagle3 capture layers (``M3DecoderLayer._feed_forward_boundary``). + +A capture layer with a next norm keeps the fused boundary (all-reduce + residual + next norm in one kernel) and +captures the reduced residual sum it writes; the last layer captures its raw output with the residual, as the +in-tree capture does. CPU only: the boundary and the capture are stand-ins. +""" + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 + layers, +) + + +class _Layer: + """The attributes the boundary reads, with a recording capture and next-norm stand-in.""" + + _fused_feed_forward_boundary = layers.M3DecoderLayer._fused_feed_forward_boundary + _feed_forward_boundary = layers.M3DecoderLayer._feed_forward_boundary + + def __init__(self, has_next_norm: bool): + self.layer_idx = 56 + self.next_layer_layernorm = object() if has_next_norm else None + self.captured = [] + self.boundaries = [] + + def capture(self, spec_metadata, layer_id, hidden_states, residual): + self.captured.append((layer_id, hidden_states, residual)) + + def _next_layer_norm(self, hidden_states, residual, fused_boundary): + self.boundaries.append(fused_boundary) + if self.next_layer_layernorm is None: + return hidden_states, residual + return ("normed", hidden_states), ("updated", hidden_states, residual) + + +def test_a_capture_layer_keeps_the_fused_boundary_and_captures_its_residual_sum() -> None: + layer = _Layer(has_next_norm=True) + fused = layer._fused_feed_forward_boundary(capture=True) + assert fused + h, r = torch.ones(2), torch.zeros(2) + out = layer._feed_forward_boundary(h, r, fused, True, None) + assert layer.boundaries == [True] + assert layer.captured == [(56, ("updated", h, r), None)] + assert out == (("normed", h), ("updated", h, r)) + + +def test_the_last_layer_captures_its_raw_output_with_the_residual() -> None: + layer = _Layer(has_next_norm=False) + fused = layer._fused_feed_forward_boundary(capture=True) + assert not fused + h, r = torch.ones(2), torch.zeros(2) + out = layer._feed_forward_boundary(h, r, fused, True, None) + assert layer.captured == [(56, h, r)] and layer.boundaries == [False] + assert out == (h, r) + + +def test_a_layer_without_capture_is_unchanged() -> None: + layer = _Layer(has_next_norm=True) + assert layer._fused_feed_forward_boundary(capture=False) + h, r = torch.ones(2), torch.zeros(2) + layer._feed_forward_boundary(h, r, True, False, None) + assert not layer.captured and layer.boundaries == [True] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_dispatch.py new file mode 100644 index 000000000000..4968e98e0177 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_dispatch.py @@ -0,0 +1,995 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tp8 decode target's dispatch table (``dispatch.py``): its row ranges against the catalog cells that certify +them and the kernels' own limits, the selection for every decode row count, and the miss / uncovered policy. + +CPU only: nothing here launches a kernel. +""" + +import ast +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 + dispatch, + layers, + linears, + spec_worker, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 + m3_attn_lc8, + m3_draft_lowm, + m3_draft_step, + m3_fused_score_select, + m3_mc_allreduce, + m3_shard_accept, + triton_sparse_decode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels.mx_gemv import ( # noqa: E501 + op as mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.modeling import ( # noqa: E501 + DecodeTarget, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention import ( # noqa: E501 + m3_fused_score_select as fused_score_select_entry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention import ( # noqa: E501 + m3_sparse_decode_lc8 as lc8_entry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm import ( # noqa: E501 + m3_draft_lowm as draft_lowm_entry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm.m3_mx_gemv import ( # noqa: E501 + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec import ( # noqa: E501 + m3_draft_step as draft_step_entry, +) +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.linear import TensorParallelMode + +_TARGET = ( + "tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8" +) + +PROJECTIONS = (QKV_SPARSE, QKV, O_PROJ, DENSE_GATE_UP, DENSE_DOWN) +ROWS = range(1, dispatch.DECODE_MAX_TOKENS + 2) # every decode row count and the first one past it + +# The Eagle3 drafter's BF16 linears at TP8 (per-rank qkv, o_proj, gate_up, down shards and the replicated fc). +DRAFTER_LINEARS = ((1280, 12288), (6144, 1024), (4608, 6144), (6144, 2304), (6144, 18432)) + + +@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() + + +class _Log: + def __init__(self): + self.info, self.warning = [], [] + + +@pytest.fixture +def log(monkeypatch) -> _Log: + """The dispatch's log lines, by level.""" + rec = _Log() + monkeypatch.setattr(dispatch.logger, "info", rec.info.append) + monkeypatch.setattr(dispatch.logger, "warning", rec.warning.append) + return rec + + +# ------------------------------------------------------------------ the ranges against their certificates + + +@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_removed_tp4_paths_are_gone() -> None: + """The tp8 target carries none of the tp4 target's coded heads, low-M MXFP8 GEMM, 9..13-row MoE engine or + attention megakernels: no fast path, range constant or kernel module of theirs is left to select.""" + names = {r.path for r in dispatch.FAST_PATHS} + assert not names & { + "m3_target_head", + "mxfp8_low_m_gemm", + "m3_moe16_engine", + "attn_mega", + "attn_c1", + "fused_decode_attn", + } + for constant in ( + "TARGET_HEAD", + "TARGET_HEAD_MAX_ROWS", + "DRAFT_HEAD_CODED", + "LOW_M_GEMM", + "LOW_M_GEMM_PATH", + "LOW_M_GEMM_MAX_TOKENS", + "MOE16_ENGINE", + "MOE16_MIN_TOKENS", + "MOE16_MAX_TOKENS", + "ATTN_MEGA", + "ATTN_C1", + "ATTN_C1_QUERY_LEN", + "FUSED_DECODE_ATTN", + "FUSED_ATTN_MAX_QUERY_LEN", + "FUSED_ATTN_MAX_TOKENS", + "SHARED_GATE_UP", + "SHARED_DOWN", + ): + assert not hasattr(dispatch, constant), constant + for module in ( + "kernels.m3_target_head", + "kernels.mxfp8_low_m_gemm", + "kernels.m3_attn_mega", + "kernels.m3_attn_mega_c1", + "kernels.m3_fused_attn", + "kernels.m3_attn_lc", + "ops.gemm.m3_target_head", + "ops.gemm.mxfp8_low_m_gemm", + "ops.moe.m3_moe16_engine", + "ops.attention.m3_attn_mega", + "ops.attention.m3_attn_c1", + "ops.attention.m3_fused_decode_attn", + "ops.attention.m3_sparse_decode_lc", + ): + assert importlib.util.find_spec(f"{_TARGET}.{module}") is None, module + for method in ("prebuild_target_head", "_m3_coded_head"): + assert not hasattr(spec_worker.M3Eagle3Worker, method), method + for method in ("moe16_engine", "_build_moe16_engine"): + assert not hasattr(layers.M3MoE, method), method + + +# ------------------------------------------------------------------ the ranges against the kernels' own limits + + +@pytest.mark.parametrize("nk", PROJECTIONS, ids=str) +def test_gemv_rows_are_the_rows_the_kernel_serves(nk) -> None: + n, k = nk + served = {m for m in range(1, mx_gemv.MAX_TOKENS + 1) if mx_gemv.serves(n, k, m, False)} + assert dispatch.MX_GEMV_ROWS[nk] == served + served_bf16 = [m for m in range(1, mx_gemv.MAX_TOKENS + 1) if mx_gemv.serves(n, k, m, True)] + assert served_bf16 == list(range(1, dispatch.MX_GEMV_BF16_MAX_TOKENS[nk] + 1)) + + +def test_the_gemv_tables_are_the_tp8_projections() -> None: + assert set(dispatch.MX_GEMV_ROWS) == set(dispatch.MX_GEMV_BF16_MAX_TOKENS) == set(PROJECTIONS) + assert PROJECTIONS == ((1536, 6144), (1280, 6144), (6144, 1024), (3072, 6144), (6144, 1536)) + for n, k in PROJECTIONS: + assert mx_gemv.get_config(n, k) is not None, (n, k) + + +@pytest.mark.parametrize("nk,push", [((768, 6144), None), ((6144, 384), "moe")], ids=str) +def test_the_shared_expert_projections_take_no_gemv(nk, push) -> None: + """At TP8 the shared expert (768 x 6144 gate_up, 6144 x 384 down) runs inside the MoE engine on decode rows and on + the generic GEMM otherwise: the GEMV has no config for it and the dispatch never routes a call there.""" + assert mx_gemv.get_config(*nk) is None + module = SimpleNamespace(weight=torch.zeros(1, 1).expand(*nk), m3_push=push, tp_size=8) + assert linears.gemv_variants(module) == [] + + +def test_the_plans_are_the_all_reduce_buckets() -> None: + assert tuple(p.max_tokens for p in dispatch.DECODE_PLANS) == m3_mc_allreduce.BUCKETS + assert dispatch.DECODE_MAX_TOKENS == m3_mc_allreduce.MAX_TOKENS == mx_gemv.MAX_TOKENS + + +def test_the_thresholds_are_the_kernels() -> None: + assert dispatch.SPARSE_DECODE_MQ_MIN_TOKENS == triton_sparse_decode.MQ_MIN_TOKENS + assert ( + dispatch.MOE_ENGINE_MAX_TOKENS < dispatch.ATTN_PUSH_MAX_TOKENS <= dispatch.DECODE_MAX_TOKENS + ) + assert dispatch.MC_FINALIZE_MIN_TOKENS == dispatch.MOE_ENGINE_MAX_TOKENS + 1 + verify_rows = dispatch.SHARD_ACCEPT_MAX_REQUESTS * (dispatch.SHARD_ACCEPT_DRAFT_LEN + 1) + assert verify_rows == dispatch.DECODE_MAX_TOKENS == m3_shard_accept.SHARD_ROWS + assert dispatch.FUSED_SCORE_SELECT_MAX_REQUESTS == fused_score_select_entry.MAX_REQUESTS == 2 + assert dispatch.INDEX_SCORE_EARLY_MAX_REQUESTS == 0 + assert dispatch.FC_EARLY_MAX_ROWS == 0 + assert dispatch.LC8_PREWAIT_MAX_ROWS == 8 + assert dict(dispatch.DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH) == {1: 16384} + assert dict(dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH) == {1: 16384} + + +def test_no_fast_path_quantizes_a_bf16_head_or_the_drafter() -> None: + """The target carries the exact-numerics configuration only: the lm_head and the Eagle3 drafter run at the + checkpoint's precision, so no fast path (nor a range constant) of an FP8 / NVFP4 / pruned head or of a quantized + drafter projection is left to select.""" + names = {r.path for r in dispatch.FAST_PATHS} + assert not names & { + "fp8_lm_head", + "draft_head", + "m3_draft_gemv", + "m3_draft_in_gemv", + "m3_gemv_ar", + } + for constant in ( + "FP8_LM_HEAD", + "FP8_LM_HEAD_MAX_TOKENS", + "DRAFT_HEAD", + "DRAFT_HEAD_MAX_TOKENS", + ): + assert not hasattr(dispatch, constant), constant + + +# ------------------------------------------------------------------ selection, every decode row count + + +def _expected_moe(rows: int): + if rows <= 8: + return "m3_moe_engine" + if rows <= 48: + return "mc_finalize_allreduce_norm" + return None + + +def _expected_router(rows: int): + if rows == 1: + return "m3_router_gemv" + if rows <= 48: + return "m3_router_tc" + return None + + +@pytest.mark.parametrize("rows", ROWS) +def test_selection(rows) -> None: + assert dispatch.is_decode_step(rows) == (rows <= 48) + if rows <= 48: + assert dispatch.plan_for(rows).name == ( + "tiny" if rows <= 8 else "small" if rows <= 16 else "verify" + ) + else: + with pytest.raises(AssertionError, match="no decode plan"): + dispatch.plan_for(rows) + assert dispatch.moe_path(rows) == _expected_moe(rows) + assert dispatch.router_path(rows) == _expected_router(rows) + assert dispatch.EMBED_FIRSTNORM.covers(rows) == (rows <= 8) + assert dispatch.ATTN_PUSH.covers(rows) == (rows <= 16) + assert dispatch.MOE_PUSH.covers(rows) == (rows <= 8) + assert dispatch.MC_ALLREDUCE.covers(rows) == (rows <= 48) + assert dispatch.FUSED_SCORE_SELECT.covers(rows) == (rows <= 2) + + +@pytest.mark.parametrize("rows", ROWS) +def test_linear_selection(rows) -> None: + bf16_max = { + QKV_SPARSE: 16, + QKV: 16, + O_PROJ: 48, + DENSE_GATE_UP: 16, + DENSE_DOWN: 32, + } + for nk in PROJECTIONS: + assert dispatch.linear_path(*nk, rows, True) == ( + "m3_mx_gemv" if rows <= 48 else "flashinfer_mxfp8_gemm" + ) + assert dispatch.linear_path(*nk, rows, False) == ( + "m3_mx_gemv_bf16" if rows <= bf16_max[nk] else "mxfp8_gemm" + ) + # A projection outside the table (the shared expert's, which runs in the MoE engine) takes the generic GEMM. + for nk in ((768, 6144), (6144, 384), (4096, 4096)): + assert dispatch.linear_path(*nk, rows, True) == "flashinfer_mxfp8_gemm" + assert dispatch.linear_path(*nk, rows, False) == "mxfp8_gemm" + + +def test_draft_step_selection() -> None: + """The fused draft loop takes the configured draft length at 1..8 requests (one head row per request, the + kernel's limit); its draft input every pass's rows: the first pass's 4 per request, then one per request.""" + assert ( + dispatch.DRAFT_STEP_MAX_REQUESTS + == draft_step_entry.MAX_TOKENS + == m3_draft_step.MAX_TOKENS + == 8 + ) + assert dispatch.DRAFT_STEP_DRAFT_LEN == dispatch.SHARD_ACCEPT_DRAFT_LEN + assert draft_step_entry.WORLD == 8 and draft_step_entry.DRAFT_SHARD * 8 == 200064 + for draft_len in range(1, 6): + for nreq in range(1, 14): + covered = draft_len == 3 and nreq <= 8 + assert dispatch.DRAFT_STEP.covers((draft_len, nreq)) == covered + assert dispatch.DRAFT_INPUT.keys == frozenset(range(1, 9)) | {12, 16, 20, 24, 28, 32} + + +def test_verify_selection() -> None: + for draft_len in range(1, 6): + for nreq in range(1, 14): + covered = draft_len == 3 and nreq * (draft_len + 1) <= 48 + assert dispatch.SHARD_ACCEPT.covers((draft_len, nreq)) == covered + + +def test_the_index_scorer_runs_its_prefetching_schedule_only() -> None: + """Spans of at most two requests score in the fused score + select; the scorer covers only the prefetching + kernel (no one-request early-issue schedule is ever selected).""" + assert dispatch.INDEX_SCORE.keys == frozenset({"prefetch"}) + assert not dispatch.INDEX_SCORE.covers("early") + for cell in dispatch.INDEX_SCORE.entry.CELLS: + assert not cell.spec["early"] + assert len(cell.spec["seq_lens"]) > dispatch.FUSED_SCORE_SELECT_MAX_REQUESTS + + +def test_the_drafter_decode_split_is_capped_at_one_request() -> None: + assert dispatch.DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH[1] == 16384 + assert 2 not in dispatch.DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH + + +# ------------------------------------------------------------------ the long-context sparse decode (lc8) + + +def test_lc8_variants_are_the_prewait_threshold() -> None: + assert dispatch.SPARSE_DECODE_LC8.keys == frozenset({"prewait", "plain"}) + for cell in lc8_entry.m3_sparse_decode_lc8.CELLS: + rows = len(cell.spec["seq_lens"]) * cell.spec["decode_query_len"] + if cell.spec["prewait"]: + assert rows <= dispatch.LC8_PREWAIT_MAX_ROWS, cell.why + + +def _lc8_operands(total_q: int): + q = torch.zeros(total_q, lc8_entry.NUM_HEADS, lc8_entry.HEAD_DIM, dtype=torch.float8_e4m3fn) + pages = torch.zeros(1, 1, lc8_entry.PAGE, lc8_entry.HEAD_DIM // 2, dtype=torch.uint8) + scales = torch.zeros(1, 1, lc8_entry.PAGE, lc8_entry.HEAD_DIM // 16, dtype=torch.uint8) + topk = torch.full((1, total_q, 16), -1, dtype=torch.int32) + out = torch.zeros(total_q, lc8_entry.NUM_HEADS, lc8_entry.HEAD_DIM, dtype=torch.bfloat16) + return q, pages, scales, topk, out + + +@pytest.mark.parametrize("total_q", [1, 4, 7, 8, 9, 16, 32, 48]) +def test_lc8_prewait_follows_the_row_threshold(monkeypatch, total_q) -> None: + """The call site stages before the grid wait exactly up to ``LC8_PREWAIT_MAX_ROWS`` query rows.""" + calls = [] + monkeypatch.setattr(m3_attn_lc8, "m3_sparse_decode_lc8", lambda *args: calls.append(args)) + q, pages, scales, topk, out = _lc8_operands(total_q) + assert m3_attn_lc8.lc8_applies(q, pages, pages, topk, scales, scales, out) + m3_attn_lc8.m3_lc8_attn( + q, + pages, + pages, + topk, + torch.zeros(1, 1, dtype=torch.int32), + torch.zeros(1, dtype=torch.int32), + sm_scale=0.1, + output=out, + decode_query_len=1, + k_block_scale=scales, + v_block_scale=scales, + k_global_scale=torch.ones(1), + v_global_scale=torch.ones(1), + ) + (args,) = calls + assert args[-1] == (total_q <= dispatch.LC8_PREWAIT_MAX_ROWS) + assert args[13] == lc8_entry.WIDTH // 128 == 8 + + +def test_lc8_takes_the_tp8_geometry_only() -> None: + q, pages, scales, topk, out = _lc8_operands(4) + assert m3_attn_lc8.lc8_applies(q, pages, pages, topk, scales, scales, out) + q16 = torch.zeros(4, 16, 128, dtype=torch.float8_e4m3fn) + out16 = torch.zeros(4, 16, 128, dtype=torch.bfloat16) + assert not m3_attn_lc8.lc8_applies(q16, pages, pages, topk, scales, scales, out16) + + +# ------------------------------------------------------------------ require mode and fallback logging + + +def test_a_miss_raises_under_require() -> None: + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_moe_engine'.* 5 rows.*not built"): + dispatch.missed("m3_moe_engine", 5, "the engine is not built") + + +def test_a_miss_logs_once_per_reason_otherwise(log) -> None: + for _ in range(3): + dispatch.missed("m3_moe_engine", 5, "the engine is not built") + dispatch.missed("m3_moe_engine", 6, "the engine failed to build") + assert len(log.warning) == 2 + assert "m3_moe_engine at 5 rows: the engine is not built" in log.warning[0] + + +def test_uncovered_never_raises_and_logs_once(log) -> None: + dispatch.set_require(True) + for _ in range(3): + dispatch.uncovered("m3_fused_score_select", "more than 2 requests: scorer + selector") + assert len(log.info) == 1 and not log.warning + + +class _Core: + """The parts of the text model ``DecodeTarget._embed`` reads: no embedding owners built.""" + + def __init__(self): + self._embed_owners = {} + + @staticmethod + def embed_tokens(input_ids: torch.Tensor) -> torch.Tensor: + return input_ids.to(torch.float32)[:, None].expand(-1, 4) + + +@pytest.mark.parametrize("rows", (1, 8)) +def test_embedding_without_its_owner_raises_under_require(rows) -> None: + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=rf"'embed_firstnorm'.* {rows} rows.*owner" + ): + DecodeTarget(_Core())._embed(torch.arange(rows, dtype=torch.int32), rows) + + +def test_embedding_without_its_owner_falls_back_and_logs_once(log) -> None: + target = DecodeTarget(_Core()) + for _ in range(2): + hidden, residual = target._embed(torch.arange(3, dtype=torch.int32), 3) + assert residual is None and hidden.shape == (3, 4) + assert len(log.warning) == 1 and "embed_firstnorm at 3 rows" in log.warning[0] + + +def test_embedding_past_its_rows_is_uncovered_in_every_mode(log) -> None: + dispatch.set_require(True) + hidden, residual = DecodeTarget(_Core())._embed(torch.arange(9, dtype=torch.int32), 9) + assert residual is None and hidden.shape == (9, 4) + assert len(log.info) == 1 and "embed_firstnorm" in log.info[0] and not log.warning + + +def test_a_push_without_its_workspace_raises_under_require(monkeypatch) -> None: + monkeypatch.setattr(m3_mc_allreduce, "get_workspace", lambda mapping, hidden: None) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'attn_push'.* 4 rows"): + DecodeTarget(_Core())._attn_push_target(_Layer(), 4, 6144) + + +def test_a_push_without_its_workspace_falls_back_and_logs_once(monkeypatch, log) -> None: + monkeypatch.setattr(m3_mc_allreduce, "get_workspace", lambda mapping, hidden: None) + target = DecodeTarget(_Core()) + assert target._attn_push_target(_Layer(), 4, 6144) is None + assert target._moe_push_targets(_Layer(), 4, 6144) is None + assert target._attn_push_target(_Layer(), 4, 6144) is None + assert len(log.warning) == 2 + assert "attn_push at 4 rows" in log.warning[0] and "moe_push at 4 rows" in log.warning[1] + + +class _Layer: + """The parts of a decoder layer the push targets read.""" + + mapping = None + + +def test_every_dispatch_name_the_target_uses_exists() -> None: + """A fallback message or selection that names a missing ``dispatch.X`` fails only when that branch runs.""" + target_dir = Path(dispatch.__file__).parent + missing = [] + for path in sorted(target_dir.rglob("*.py")): + for node in ast.walk(ast.parse(path.read_text())): + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "dispatch" + and not hasattr(dispatch, node.attr) + ): + missing.append( + f"{path.relative_to(target_dir)}:{node.lineno}: dispatch.{node.attr}" + ) + assert not missing, missing + + +# ------------------------------------------------------------------ the fused index score + top-16 select + + +class _FssConfig: + num_index_heads = 1 + sparse_index_dim = 128 + num_kv_heads = 1 + init_blocks = 0 + local_blocks = 1 + + +def _fss_layer(): + """What ``M3Attention._fused_score_select`` reads of its layer.""" + return SimpleNamespace(attn=SimpleNamespace(m3_config=_FssConfig()), layer_idx=5) + + +def _fss_metadata(rows: int, qlen: int, max_tiles: int = 8192, n_valid: bool = True): + """A decode span's metadata as the fused score + select reads it (CPU stand-ins, never launched).""" + return SimpleNamespace( + msa_idx_k_cache=lambda layer_idx: torch.zeros(4, 1, 128, 128, dtype=torch.float8_e4m3fn), + msa_worst_case_max_k_tiles=max_tiles, + msa_n_valid_blocks=torch.ones(rows * qlen, dtype=torch.int32) if n_valid else None, + msa_block_table=torch.zeros(rows, 4, dtype=torch.int32), + msa_seq_lens_cuda=torch.full((rows,), 300, dtype=torch.int32), + ) + + +def _fss(rows: int, qlen: int = 4, dtype=torch.float8_e4m3fn, **md): + idx_q = torch.zeros(rows * qlen, 128, dtype=dtype) + return layers.M3Attention._fused_score_select( + _fss_layer(), idx_q, _fss_metadata(rows, qlen, **md), rows, qlen, rows * qlen + ) + + +def test_fused_score_select_covers_one_and_two_requests() -> None: + assert dispatch.FUSED_SCORE_SELECT.keys == frozenset({1, 2}) + assert {c.spec["batch"] for c in dispatch.FUSED_SCORE_SELECT.entry.CELLS} == {1, 2} + + +@pytest.mark.parametrize( + "kwargs,reason", + [ + (dict(dtype=torch.bfloat16), "not FP8"), + (dict(qlen=5), "not the kernel's geometry"), + (dict(max_tiles=16384), "do not fit the stage ring"), + (dict(n_valid=False), "not one decode span"), + ], + ids=["bf16_index", "qlen5", "ring", "no_valid_counts"], +) +@pytest.mark.parametrize("rows", [1, 2]) +def test_a_covered_span_failing_a_precondition_raises_under_require(rows, kwargs, reason) -> None: + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=rf"'m3_fused_score_select'.* {rows} rows.*{reason}" + ): + _fss(rows, **kwargs) + + +def test_a_covered_span_failing_a_precondition_falls_back_and_logs_once(log) -> None: + for _ in range(2): + assert _fss(2, dtype=torch.bfloat16) is None + assert len(log.warning) == 1 and "m3_fused_score_select at 2 rows" in log.warning[0] + + +def test_a_capture_without_the_select_scratch_raises_under_require(monkeypatch) -> None: + monkeypatch.setattr(m3_fused_score_select, "buffers", lambda device, max_tiles: None) + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=r"'m3_fused_score_select'.* 1 rows.*scratch.*captured" + ): + _fss(1) + + +@pytest.mark.parametrize("rows", [3, 8]) +def test_a_span_past_two_requests_is_uncovered_in_every_mode(rows, log) -> None: + dispatch.set_require(True) + assert _fss(rows) is None + assert len(log.info) == 1 and "m3_fused_score_select" in log.info[0] and not log.warning + + +# ------------------------------------------------------------------ the drafter's low-M linears + + +def test_draft_lowm_keys_are_the_tactic_rows() -> None: + tactics = draft_lowm_entry.TACTICS + assert set(tactics) == set(DRAFTER_LINEARS) == set(dispatch.DRAFT_LOWM_ROWS) + expected = frozenset( + (nk, m) for nk, table in tactics.items() for m in range(1, table[-1][0] + 1) + ) + assert dispatch.DRAFT_LOWM.keys == expected + for nk, table in tactics.items(): + rows = [entry[0] for entry in table] + assert rows == sorted(rows) and len(set(rows)) == len(rows), nk + for m in range(1, table[-1][0] + 1): + assert draft_lowm_entry.tactic_for(*nk, m) is not None, (nk, m) + assert draft_lowm_entry.tactic_for(*nk, table[-1][0] + 1) is None, nk + + +def _lowm_operands(nk, m: int): + """An [N, K] BF16 weight (a broadcast view: the dispatch reads its shape and dtype only) and an [m, K] input.""" + n, k = nk + weight = torch.zeros(1, 1, dtype=torch.bfloat16).expand(n, k) + return weight, torch.zeros(m, k, dtype=torch.bfloat16) + + +def _draft_linear(weight, x, bias=None): + return m3_draft_lowm.draft_linear(weight, x, bias) + + +@pytest.mark.parametrize("nk", DRAFTER_LINEARS, ids=str) +def test_draft_linear_without_its_prebuilt_kernel_raises_under_require(monkeypatch, nk) -> None: + monkeypatch.setattr(m3_draft_lowm, "_BUILT", set()) + weight, x = _lowm_operands(nk, 1) + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=r"'m3_draft_lowm'.* 1 rows.*not compiled before warmup" + ): + _draft_linear(weight, x) + + +def test_draft_linear_without_its_prebuilt_kernel_falls_back_and_logs_once( + monkeypatch, log +) -> None: + monkeypatch.setattr(m3_draft_lowm, "_BUILT", set()) + weight, x = _lowm_operands((6144, 1024), 2) + for _ in range(2): + assert _draft_linear(weight, x) is None + assert len(log.warning) == 1 and "m3_draft_lowm at 2 rows" in log.warning[0] + + +def test_draft_linear_on_a_non_bf16_activation_is_a_miss() -> None: + weight, x = _lowm_operands((6144, 1024), 1) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_draft_lowm'.*not BF16"): + _draft_linear(weight, x.float()) + + +def test_draft_linear_on_a_strided_activation_keeps_the_in_tree_gemm(log) -> None: + """The first pass's fc may read a column window of the captured hidden states: uncovered, never raised.""" + weight, x = _lowm_operands((6144, 1024), 2) + wide = x.new_zeros(2, 2048) + wide[:, :1024] = x + assert not wide[:, :1024].is_contiguous() + dispatch.set_require(True) + assert _draft_linear(weight, wide[:, :1024]) is None + assert len(log.info) == 1 and not log.warning + + +def test_draft_linear_past_its_rows_is_uncovered_in_every_mode(monkeypatch, log) -> None: + monkeypatch.setattr(m3_draft_lowm, "_BUILT", set()) + dispatch.set_require(True) + weight, x = _lowm_operands((6144, 1024), 3) # the o_proj table ends at 2 rows + assert _draft_linear(weight, x) is None + weight, x = _lowm_operands((1280, 12288), 9) + assert _draft_linear(weight, x) is None + weight, x = _lowm_operands((4096, 1024), 1) # not a drafter shape + assert _draft_linear(weight, x) is None + assert len(log.info) == 3 and not log.warning + + +def test_draft_linear_with_a_bias_is_not_taken(monkeypatch, log) -> None: + monkeypatch.setattr(m3_draft_lowm, "_BUILT", set()) + dispatch.set_require(True) + weight, x = _lowm_operands((6144, 1024), 1) + assert _draft_linear(weight, x, torch.zeros(6144, dtype=torch.bfloat16)) is None + assert _draft_linear(weight, x[None]) is None + assert not log.info and not log.warning + + +def test_prebuild_drafter_linears_caps_the_drafter_attention(monkeypatch) -> None: + """The drafter attention gets ``DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH`` and every drafter linear goes to the + low-M method before warmup; without a drafter nothing is built.""" + installed = [] + monkeypatch.setattr(m3_draft_lowm, "install", lambda dm: installed.append(dm) or 5) + monkeypatch.setattr(m3_draft_lowm, "prebuild", lambda dm: 17) + worker = spec_worker.M3Eagle3Worker.__new__(spec_worker.M3Eagle3Worker) + assert worker.prebuild_drafter_linears(None) == 0 and not installed + attn = SimpleNamespace() + drafter = SimpleNamespace( + model=SimpleNamespace(midlayer=SimpleNamespace(self_attn=SimpleNamespace(attn=attn))) + ) + assert worker.prebuild_drafter_linears(drafter) == 17 + assert installed == [drafter] + assert attn.dense_decode_max_seq_len_by_batch is dispatch.DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH + + +# ------------------------------------------------------------------ the fused draft loop's preconditions + +_HEAD_ROWS = 64 +_HIDDEN = 32 + + +def _drafter(head_weight=None): + """A drafter with only what M3Eagle3Worker._m3_drafter_supported reads: one plain Eagle3 layer with a gated MLP, + a whole-table embedding, all-reduced o_proj and down_proj and a vocab-sharded BF16 head (TP8).""" + mlp = GatedMLP.__new__(GatedMLP) + torch.nn.Module.__init__(mlp) + mlp.down_proj = SimpleNamespace(all_reduce=object()) + layer = SimpleNamespace( + _next_layer_regular=False, + mlp=mlp, + self_attn=SimpleNamespace(o_proj=SimpleNamespace(all_reduce=object())), + ) + model = SimpleNamespace( + num_layers=1, + _use_mla=False, + _eh_proj_before_attn=False, + midlayer=layer, + embed_tokens=SimpleNamespace(tp_size=1), + ) + weight = ( + head_weight + if head_weight is not None + else torch.zeros((_HEAD_ROWS, _HIDDEN), dtype=torch.bfloat16) + ) + lm_head = SimpleNamespace( + weight=weight, bias=None, tp_mode=TensorParallelMode.COLUMN, padding_size=0, tp_size=8 + ) + return SimpleNamespace(model=model, lm_head=lm_head) + + +def _worker(): + """An M3Eagle3Worker past the step's verify: the fused verify left the first-pass gather rows and the fused + first draft inputs; nothing else of the in-tree worker is read.""" + worker = spec_worker.M3Eagle3Worker.__new__(spec_worker.M3Eagle3Worker) + worker.mapping = SimpleNamespace(tp_size=8, tp_rank=1, enable_attention_dp=False) + worker._d2t = None + worker._m3_draft_head = None + worker._m3_first_gather_ids = torch.zeros(1, dtype=torch.int64) + worker._m3_fused_first_inputs = True + return worker + + +_ATTN_MD = SimpleNamespace(kv_lens_cuda=None) +_SPEC_MD = SimpleNamespace(runtime_draft_len=3) + + +@pytest.fixture +def no_gpu_draft_head(monkeypatch): + """The fused head's CUDA-only lookups replaced: a 148-SM device, and a built exchange workspace.""" + monkeypatch.setattr(m3_draft_step, "_num_sms", lambda: 148) + monkeypatch.setattr(m3_draft_step, "_WS", object()) + + +def test_draft_loop_without_the_prebuilt_head_raises_under_require(no_gpu_draft_head) -> None: + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=r"'m3_draft_step'.* 4 rows.*prebuild_draft_head" + ): + _worker()._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, _drafter()) + + +def test_draft_loop_without_the_prebuilt_head_falls_back_and_logs_once( + no_gpu_draft_head, log +) -> None: + worker = _worker() + for _ in range(2): + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, _drafter()) + assert len(log.warning) == 1 and "m3_draft_step at 4 rows" in log.warning[0] + + +def test_draft_loop_with_a_head_built_for_another_weight_raises_under_require( + no_gpu_draft_head, +) -> None: + worker = _worker() + assert worker.prebuild_draft_head(_drafter()) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_draft_step'"): + worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 2, _drafter()) + + +def test_draft_loop_takes_a_covered_step_once_the_head_is_built(no_gpu_draft_head) -> None: + dispatch.set_require(True) + drafter = _drafter() + worker = _worker() + assert worker.prebuild_draft_head(drafter) + head = worker._m3_draft_head + assert head.weight is drafter.lm_head.weight and head.vocab_offset == _HEAD_ROWS + assert torch.equal(draft_step_entry.from_fragments(head.wf), drafter.lm_head.weight) + for nreq in range(1, 9): + assert worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, nreq, drafter) + + +def test_draft_loop_streams_the_bf16_fragments_only(no_gpu_draft_head) -> None: + """The tp8 draft head is the fragment-ordered BF16 copy at every K (no coded copy is built, whatever K), it + decodes to the weight bit for bit, and the op refuses a coded stream.""" + dispatch.set_require(True) + g = torch.Generator().manual_seed(11) + weight = (torch.randn((_HEAD_ROWS, 128), generator=g) * 0.02).to(torch.bfloat16) + drafter = _drafter(weight) + worker = _worker() + assert worker.prebuild_draft_head(drafter) + head = worker._m3_draft_head + assert not hasattr(head, "coded") and not hasattr(head, "exc_off") + assert head.weight is weight and head.vocab_offset == _HEAD_ROWS + decoded = draft_step_entry.head_weight(head.wf, 128) + assert torch.equal(decoded.view(torch.int16), weight.view(torch.int16)) + with pytest.raises(NotImplementedError, match="BF16 fragments only"): + draft_step_entry.head_weight(head.wf, 128, torch.zeros(5, dtype=torch.int32)) + assert all(not c.spec["coded"] for c in draft_step_entry.m3_draft_step.CELLS) + for nreq in range(1, 9): + assert worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, nreq, drafter) + + +def test_draft_loop_past_its_requests_is_uncovered_in_every_mode(no_gpu_draft_head, log) -> None: + dispatch.set_require(True) + drafter = _drafter() + worker = _worker() + assert worker.prebuild_draft_head(drafter) + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 9, drafter) + assert not worker._m3_fused_loop_ok(_ATTN_MD, SimpleNamespace(runtime_draft_len=2), 4, drafter) + assert len(log.info) == 2 and not log.warning + + +def test_draft_loop_needs_the_fused_first_inputs(no_gpu_draft_head, log) -> None: + """A step whose first draft inputs came from the in-tree path (or that the fused verify did not prepare) drafts + with the in-tree loop, silently: it is not a draft-step shape.""" + dispatch.set_require(True) + drafter = _drafter() + worker = _worker() + assert worker.prebuild_draft_head(drafter) + worker._m3_fused_first_inputs = False + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + worker._m3_fused_first_inputs, worker._m3_first_gather_ids = True, None + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + assert not log.info and not log.warning + + +def test_an_unsupported_drafter_is_uncovered(no_gpu_draft_head, log) -> None: + dispatch.set_require(True) + worker = _worker() + padded = _drafter() + padded.lm_head.padding_size = 16 + two_layers = _drafter() + two_layers.model.num_layers = 2 + for drafter in (padded, two_layers): + assert not worker.prebuild_draft_head(drafter) + assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + assert log.info and not log.warning + + +def test_a_capture_without_the_exchange_workspace_raises_under_require(monkeypatch) -> None: + monkeypatch.setattr(m3_draft_step, "_num_sms", lambda: 148) + monkeypatch.setattr(m3_draft_step, "_WS", None) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + drafter = _drafter() + worker = _worker() + assert worker.prebuild_draft_head(drafter) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_draft_step'.*workspace.*captured"): + worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) + + +# ------------------------------------------------------------------ the fc fork (never early at tp8) + + +class _CaptureMetadata: + def __init__(self, nreq: int): + self.layers_to_capture = (1, 29, 56) + self.num_capture_tokens = nreq * 4 + self.runtime_draft_len = 3 + self.captured = [] + + def maybe_capture_hidden_states(self, layer_id, hidden_states, residual) -> None: + self.captured.append(layer_id) + + +def _fc_early_worker(monkeypatch, eligible: bool = True): + worker = spec_worker.M3Eagle3Worker.__new__(spec_worker.M3Eagle3Worker) + worker._m3_capture_norm = None + worker._m3_fc_norm_w = None + worker._m3_fc_early_drafter = None + worker._m3_fc_early_stream = None + forks = [] + monkeypatch.setattr(worker, "tail_eligible", lambda *args: eligible) + + def fork(md, num_tokens, draft_model): + forks.append((num_tokens, draft_model)) + return f"stream-{len(forks)}" + + monkeypatch.setattr(worker, "_m3_fc_fork", fork) + return worker, forks + + +@pytest.mark.parametrize("nreq,eligible", [(1, True), (2, True), (3, True), (8, True), (1, False)]) +def test_no_early_fork_at_tp8(monkeypatch, nreq, eligible) -> None: + """``FC_EARLY_MAX_ROWS`` is 0: no verify (not even one request's 4 rows) forks the drafter's fc at the last + captured layer; every captured layer still captures and the fc forks beside the verify head instead.""" + worker, forks = _fc_early_worker(monkeypatch, eligible) + md = _CaptureMetadata(nreq) + worker._m3_fc_early_stream = "stale" + worker.arm_capture_norm(SimpleNamespace(num_seqs=nreq), md, object(), object()) + assert worker._m3_fc_early_drafter is None + for layer_id in md.layers_to_capture: + worker.capture_hidden_states(md, layer_id, None) + worker.capture_hidden_states(md, spec_worker.FINAL_HIDDEN_LAYER_ID, None) + assert md.captured == [1, 29, 56, spec_worker.FINAL_HIDDEN_LAYER_ID] + assert not forks and worker._m3_fc_early_stream is None + + +def test_no_early_fork_with_a_final_state_capture(monkeypatch) -> None: + """The final-state column is written after the model returns, so captures that include it keep the fc at its + stock place instead of forking at the last captured decoder layer.""" + worker, forks = _fc_early_worker(monkeypatch) + md = _CaptureMetadata(1) + md.layers_to_capture = (spec_worker.FINAL_HIDDEN_LAYER_ID, 29, 56) + worker.arm_capture_norm(SimpleNamespace(num_seqs=1), md, object(), object()) + for layer_id in (29, 56, spec_worker.FINAL_HIDDEN_LAYER_ID): + worker.capture_hidden_states(md, layer_id, None) + assert not forks and worker._m3_fc_early_stream is None + + +# ------------------------------------------------------------------ the NVFP4 producer's coverage + + +def test_the_producer_covers_the_tp8_nvfp4_layers_only() -> None: + """``dispatch.PRODUCER`` covers a sparse layer on an NVFP4 cache at the TP8 heads (8 query, one KV and one + index head); elsewhere the built-in producers run without a miss.""" + covered = layers.producer_pw_covered + assert covered(True, 8, 1, 1) + assert not covered(False, 8, 1, 1) + assert not covered(True, 16, 1, 1) + assert not covered(True, 8, 2, 2) + assert dispatch.PRODUCER.keys == frozenset({"nvfp4_tp8"}) + assert dispatch.PRODUCER.covers("nvfp4_tp8") + assert not dispatch.PRODUCER.covers("nvfp4_tp4") + + +class _FakeMsaAttn: + indexer_kv_dtype = "fp8" + + +def _producer_layer(monkeypatch, nvfp4: bool = True, heads: int = 8): + """What the producer dispatch reads before its first precondition: a fused-projection MSA sparse layer.""" + monkeypatch.setattr(layers, "MiniMaxM3MsaSparseAttention", _FakeMsaAttn) + return SimpleNamespace( + enable_fused_qkv_index_projection=True, + attn=_FakeMsaAttn(), + main_kv_is_fp8=not nvfp4, + main_kv_is_nvfp4=nvfp4, + num_heads=heads, + num_key_value_heads=1, + sparse_num_index_heads=1, + pos_embd_params=None, + ) + + +def _produce(layer): + # The packed TP8 projection: 8 query heads, K, V, index Q and index K of 128 each. + packed = torch.zeros(4, 12 * 128, dtype=torch.bfloat16) + return layers.M3Attention._fused_fp8_qkv_indexer_norm_rope_kv_insert( + layer, packed, None, SimpleNamespace() + ) + + +def test_a_covered_producer_layer_missing_a_precondition_raises_under_require(monkeypatch) -> None: + """A TP8 NVFP4 sparse layer whose step lacks a producer precondition (here the position ids) is a miss.""" + dispatch.set_require(True) + with pytest.raises( + dispatch.FastPathUnavailable, match=r"'m3_producer' covers this step at 4 rows" + ): + _produce(_producer_layer(monkeypatch)) + + +def test_a_covered_producer_layer_without_the_fused_projection_raises_under_require( + monkeypatch, +) -> None: + layer = _producer_layer(monkeypatch) + layer.enable_fused_qkv_index_projection = False + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="fuse_qkv_index_projection is disabled"): + layers.M3Attention._check_producer_projection(layer) + + +def test_an_uncovered_layer_without_the_fused_projection_is_not_a_miss(monkeypatch) -> None: + layer = _producer_layer(monkeypatch, nvfp4=False) + layer.enable_fused_qkv_index_projection = False + dispatch.set_require(True) + layers.M3Attention._check_producer_projection(layer) + + +def test_a_covered_producer_miss_falls_back_and_logs_once(monkeypatch, log) -> None: + layer = _producer_layer(monkeypatch) + assert _produce(layer) is None and _produce(layer) is None + assert len(log.warning) == 1 and "m3_producer at 4 rows" in log.warning[0] + + +@pytest.mark.parametrize("nvfp4,heads", [(False, 8), (True, 16)], ids=["fp8_cache", "tp4_heads"]) +def test_an_uncovered_producer_layer_falls_back_silently(monkeypatch, log, nvfp4, heads) -> None: + dispatch.set_require(True) + assert _produce(_producer_layer(monkeypatch, nvfp4, heads)) is None + assert not log.warning + + +def test_a_covered_producer_layer_with_the_wrong_pool_layout_is_not_supported() -> None: + """The pool-layout check the covered branch takes before the native fallback (then a miss).""" + pool = torch.zeros(4, 2, 1, 128, 32, dtype=torch.uint8) + scales = torch.zeros(4, 2, 1, 128, 8, dtype=torch.uint8) + assert not layers.producer_pw_supported(8, 1, 1, pool, scales) + good = torch.zeros(4, 2, 1, 128, 64, dtype=torch.uint8) + assert layers.producer_pw_supported(8, 1, 1, good, scales) + assert not layers.producer_pw_supported(16, 1, 1, good, scales) + assert not layers.producer_pw_supported( + 8, 1, 1, good, scales.transpose(3, 4).contiguous().transpose(3, 4) + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_prebuild.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_prebuild.py new file mode 100644 index 000000000000..580f419a3f7f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_prebuild.py @@ -0,0 +1,305 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode kernels ``build_targets`` builds before the executor's warmup (``prebuild_kernels``), and that a forward +after it never compiles one: the MoE engine of every ``MOE_ENGINE`` token count, the decode GEMVs with their push +variants, and the index scorer's prefetching kernel at every decode query length (the tp8 target never selects the +scorer's early-issue schedule: ``INDEX_SCORE_EARLY_MAX_REQUESTS`` is 0). + +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 +from unittest import mock + +import pytest +import torch +from torch import nn + +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_103__tp8 import ( # noqa: E501 + dispatch, + layers, + linears, + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels.mx_gemv import ( # noqa: E501 + op as mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.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_103__tp8.ops.gemm.m3_mx_gemv import ( # noqa: E501 + DENSE_DOWN, + O_PROJ, + QKV_SPARSE, +) + +assert torch.cuda.is_available(), "the prebuild checks require a CUDA device" + +INDEX_HEADS = 1 # index heads per rank at TP8 +INDEX_DIM = 128 +DRAFT_LEN = 3 + + +@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 _moe(layer_idx: int, shared: bool = True) -> layers.M3MoE: + """An M3MoE built by its own constructor over a stub stock block: only what prebuild_engines and the lookups + read (the builds are stubbed).""" + + def stock_init(self, model_config, aux_stream_dict, layer_idx=None): + nn.Module.__init__(self) + # One rank: the build agreement takes no collective. + self.mapping = SimpleNamespace(tp_size=1, enable_attention_dp=False) + self.shared_experts = None + + with mock.patch.object(layers.MiniMaxM3MoE, "__init__", stock_init): + moe = layers.M3MoE(None, {}, layer_idx) + moe.shared_experts = object() if shared else None + return moe + + +def _model(monkeypatch, num_layers: int = 3, indexer_kv_dtype: str = "fp8"): + """The parts of MiniMaxM3Tp8ModelingV2 prebuild_kernels reads: layers (MoE + sparse attention), the embedding + device and the speculative config.""" + monkeypatch.setattr(layers, "MiniMaxM3MsaSparseAttention", _FakeSparseAttn) + model_layers = [] + for idx in range(num_layers): + model_layers.append( + SimpleNamespace( + block_sparse_moe=_moe(idx, shared=idx != 1) if idx else None, + self_attn=SimpleNamespace( + is_sparse_attention_layer=idx != 0, attn=_FakeSparseAttn(indexer_kv_dtype) + ), + ) + ) + stub = SimpleNamespace( + layers=model_layers, + embed_tokens=SimpleNamespace(weight=torch.empty(1, device="cuda")), + model_config=SimpleNamespace( + spec_config=SimpleNamespace(max_draft_len=DRAFT_LEN, max_total_draft_tokens=DRAFT_LEN) + ), + ) + stub.decode_query_lens = lambda: modeling.MiniMaxM3Tp8ModelingV2.decode_query_lens(stub) + return stub + + +def _stub_engine_builds(monkeypatch, built: list) -> None: + def build(self, num_tokens, shared=True): + built.append(("moe", self.layer_idx, num_tokens, shared)) + return ("engine", self.layer_idx, num_tokens) + + monkeypatch.setattr(layers.M3MoE, "_build_engine", build) + + +def _forbid_compiles(monkeypatch) -> 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(layers.M3MoE, "_build_engine", boom) + monkeypatch.setattr(layers._scorer_runner(), "_compile", boom) + + +def test_decode_query_lens_follow_the_speculative_config(monkeypatch) -> None: + model = _model(monkeypatch) + assert model.decode_query_lens() == [1, 1 + DRAFT_LEN] + model.model_config.spec_config = None + assert model.decode_query_lens() == [1] + model.model_config.spec_config = SimpleNamespace( + max_draft_len=4, max_total_draft_tokens=4, draft_len_schedule={1: 4, 8: 2} + ) + assert model.decode_query_lens() == [1, 3, 5] + + +def test_prebuild_builds_every_engine_count_and_scorer_variant(monkeypatch) -> None: + built = [] + _stub_engine_builds(monkeypatch, built) + monkeypatch.setattr(modeling, "adopt_m3_linears", lambda root: built.append(("linears",))) + model = _model(monkeypatch) + times = modeling.MiniMaxM3Tp8ModelingV2.prebuild_kernels(model) + + assert built[0] == ("linears",) + assert mx_gemv.frozen(), "prebuild must close the GEMV compiles" + # Layer 0 is dense, layer 1 has no shared expert (no engine takes it); layer 2 gets every count, in order. + assert built[1:] == [("moe", 2, n, True) for n in range(1, dispatch.MOE_ENGINE_MAX_TOKENS + 1)] + moe = model.layers[2].block_sparse_moe + assert set(moe._engines) == {(n, True, True) for n in dispatch.MOE_ENGINE.keys} + assert not hasattr(moe, "_moe16_engines") + + runner = layers._scorer_runner() + want = { + (torch.float8_e4m3fn, INDEX_HEADS, qlen, INDEX_DIM, layers.INDEX_SCORE_SPLIT_K, False) + for qlen in (1, 1 + DRAFT_LEN) + } + assert set(times["scorer_keys"]) == want + assert want <= set(runner.kernel_cache) + for key in ("gemv_s", "moe_s", "scorer_s"): + assert times[key] >= 0.0 + + # A second build_targets builds nothing again. + built.clear() + modeling.MiniMaxM3Tp8ModelingV2.prebuild_kernels(model) + assert built == [("linears",)] + + +def test_engine_lookups_after_prebuild_never_build(monkeypatch) -> None: + built = [] + _stub_engine_builds(monkeypatch, built) + moe = _moe(5) + moe.prebuild_engines() + _forbid_compiles(monkeypatch) + for require in (False, True): + dispatch.set_require(require) + for n in sorted(dispatch.MOE_ENGINE.keys): + assert moe.engine(n, shared=True) == ("engine", 5, n) + + +def test_missing_engine_is_a_miss_not_a_build(monkeypatch) -> None: + _forbid_compiles(monkeypatch) + moe = _moe(7) + dispatch.set_require(False) + assert moe.engine(3, shared=True) is None + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="not built before warmup"): + moe.engine(3, shared=True) + + +def test_failed_engine_build_is_logged_then_missed(monkeypatch) -> None: + def fail(self, num_tokens, shared=True): + raise ValueError("experts backend is not TRTLLM-Gen NVFP4 SwiGLU-OAI") + + monkeypatch.setattr(layers.M3MoE, "_build_engine", fail) + moe = _moe(9) + moe.prebuild_engines() + assert all(v is None for v in moe._engines.values()) and len(moe._engines) == 8 + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="build failed"): + moe.engine(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, m3_push=push, tp_size=8) + + +@pytest.mark.parametrize( + "nk,push", + ((O_PROJ, linears.PUSH_ATTN), (QKV_SPARSE, None), (DENSE_DOWN, None)), + ids=("o_proj", "qkv_sparse", "dense_down"), +) +def test_gemvs_after_prebuild_never_compile(monkeypatch, nk, push) -> None: + use_pdl = True + module = _projection(nk, push) + variants = linears.gemv_variants(module) + 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 == linears.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) + assert linears.prebuild_m3_gemvs(module, use_pdl) > 0 + 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(DENSE_DOWN, None) + n, k = DENSE_DOWN + 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, k, device="cuda").to(torch.bfloat16), True + ) + key = mx_gemv._resolve(3, n, k, 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, k), 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: + model = _model(monkeypatch, indexer_kv_dtype=indexer_kv_dtype) + keys = layers.prebuild_index_scorer( + model.layers, model.decode_query_lens(), torch.device("cuda") + ) + # The decode query lengths 1 and 1 + DRAFT_LEN, the prefetching kernel only. + assert len(keys) == 2 and not any(key[-1] for key in keys) + _forbid_compiles(monkeypatch) + 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() + + +def test_moe_keeps_its_layer_index(): + # prebuild_engines names the layer in its build and agreement messages. + assert _moe(5).layer_idx == 5 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py new file mode 100644 index 000000000000..f0b7fcd51ecf --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py @@ -0,0 +1,199 @@ +# 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 torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 + dispatch, + layers, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 + m3_mc_allreduce as mc, +) + + +@pytest.fixture(autouse=True) +def _restore_modes(): + require = dispatch.require_mode() + yield + dispatch.set_require(require) + dispatch.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 _moe(layer_idx: int, fail: bool, world: int = 2) -> layers.M3MoE: + """An M3MoE on ``world`` simulated TP ranks whose engine builds are stubbed (and fail where ``fail``).""" + + def stock_init(self, model_config, aux_stream_dict, layer_idx=None): + nn.Module.__init__(self) + self.mapping = _mapping(world) + self.shared_experts = None + + with mock.patch.object(layers.MiniMaxM3MoE, "__init__", stock_init): + moe = layers.M3MoE(None, {}, layer_idx) + moe.shared_experts = object() + + def build(num_tokens, shared=True): + if fail: + raise RuntimeError("compile failed on this rank") + return ("engine", num_tokens) + + moe._build_engine = build + return moe + + +@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) + moes = [_moe(3, rank in failing, world) for rank in range(world)] + _run_ranks([m.prebuild_engines for m in moes]) + keys = sorted(moes[0]._engines) + assert keys and comm.calls == world * len(keys) + assert not any(hasattr(m, "_moe16_engines") for m in moes) + for m in moes[1:]: + assert m._engines.keys() == moes[0]._engines.keys() + for key in keys: + for m in moes: + assert (m._engines[key] is None) == bool(failing), key + # Every rank then takes the same path in the forward: the engine, or (after a miss) the built-in MoE. + dispatch.set_require(False) + for n in sorted(dispatch.MOE_ENGINE.keys): + got = [m.engine(n, shared=True) is None for m in moes] + assert len(set(got)) == 1, (n, got) + + +@pytest.mark.parametrize("capturing", [False, True]) +def test_fc_join_records_output_on_main(monkeypatch, capturing): + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.spec_worker import ( # noqa: E501 + 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_103__tp8/test_modeling_v2_m3_tp4_drift.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py new file mode 100644 index 000000000000..8fa2ee6f331f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tp8 target's files that are byte copies of the tp4 target's (and of the shared CuTe DSL GEMV kernel) stay so. + +The two MiniMax-M3 targets share no code by design (``test_modeling_v2_claims.test_targets_do_not_share_files``): the +tp8 target was made as a copy of the tp4 one and keeps its own copy of every file, tp8-specific or not. The files +listed here are still identical, so a fix to one copy applies to the other unchanged. This test fails on the first +edit to either side: a change to one copy must be considered for the other -- applied to both, or the file dropped +from this list with the reason the copies now differ. + +CPU only: it reads the files by path relative to the packages. +""" + +from pathlib import Path + +import pytest + +import tensorrt_llm._torch as _torch_pkg +import tensorrt_llm._torch._experimental.modeling_v2 as _modeling_v2 + +# Nothing of either target is imported (no GPU, no built extensions needed): the paths are the packages', resolved +# from the imported parents rather than from this file, which lives under tests/. +pytestmark = pytest.mark.cpu_only + +_M3 = Path(_modeling_v2.__file__).resolve().parent / "models" / "minimax_m3" +_TP4 = _M3 / "minimax_m3_nvfp4__sm_103__tp4" +_TP8 = _M3 / "minimax_m3_nvfp4__sm_103__tp8" + +#: Paths relative to the target directory that are byte copies of the tp4 target's. +IDENTICAL_TO_TP4 = ( + "kernels/long_select.py", + "kernels/m3_index_score.py", + "kernels/m3_prequant.py", + "kernels/m3_router_gemv.py", + "kernels/m3_sampled_accept.py", + "kernels/m3_shard_accept.py", + "kernels/m3_spec_tail.py", + "kernels/msa_kv_update.py", + "kernels/triton_sparse_decode_mq.py", + "ops/__init__.py", + "ops/attention/__init__.py", + "ops/attention/msa_kv_update.py", + "ops/comm/__init__.py", + "ops/gemm/__init__.py", + "ops/moe/__init__.py", + "ops/moe/m3_router_gemv.py", + "ops/moe/m3_router_gemv_fp4.py", + "ops/moe/m3_router_tc.py", + "ops/norm/__init__.py", + "ops/norm/m3_embedding_firstnorm.py", + "ops/spec/__init__.py", + "ops/spec/m3_capture_norm.py", + "ops/spec/m3_chunk_norm.py", + "ops/spec/m3_draft_input.py", + "ops/spec/m3_slot_scatter.py", + "ops/spec/m3_step_gather.py", +) + +#: The tp8 target's own copy of the CuTe DSL decode GEMV kernel (its op.py carries the TP8 tables), against the shared +#: one the tp4 target uses. +GEMV_KERNEL = "kernels/mx_gemv/m3_mx_gemv_kernel.py" +_SHARED_GEMV_KERNEL = ( + Path(_torch_pkg.__file__).resolve().parent + / "cute_dsl_kernels" + / "minimax_m3" + / "mx_gemv" + / "m3_mx_gemv_kernel.py" +) + + +def test_the_targets_are_where_this_test_reads_them() -> None: + for path in ( + _TP4 / "dispatch.py", + _TP8 / "dispatch.py", + _TP8 / GEMV_KERNEL, + _SHARED_GEMV_KERNEL, + ): + assert path.is_file(), path + + +def _differs(a: Path, b: Path) -> str: + """Empty when the two files are byte-identical, else where they first differ.""" + x, y = a.read_bytes(), b.read_bytes() + if x == y: + return "" + first = next((i for i, (p, q) in enumerate(zip(x, y)) if p != q), min(len(x), len(y))) + line = x[:first].count(b"\n") + 1 + return f"{a} and {b} differ from line {line} ({len(x)} vs {len(y)} bytes)" + + +@pytest.mark.parametrize("relative", IDENTICAL_TO_TP4) +def test_the_copy_is_the_tp4_file(relative) -> None: + ours, theirs = _TP8 / relative, _TP4 / relative + assert ours.is_file(), f"{relative}: gone from the tp8 target; drop it from IDENTICAL_TO_TP4" + assert theirs.is_file(), f"{relative}: gone from the tp4 target; drop it from IDENTICAL_TO_TP4" + diff = _differs(ours, theirs) + assert not diff, ( + f"{diff}. The targets share no code, so a change to one copy must be considered for the other: apply it to " + "both, or drop the file from IDENTICAL_TO_TP4 stating why the copies now differ." + ) + + +def test_the_gemv_kernel_is_the_shared_kernel() -> None: + diff = _differs(_TP8 / GEMV_KERNEL, _SHARED_GEMV_KERNEL) + assert not diff, ( + f"{diff}. The tp8 target's decode GEMV kernel is a copy of the shared CuTe DSL one: a change to one copy " + "must be considered for the other." + ) + + +def test_the_list_names_every_identical_file() -> None: + """A file that became identical again (or was copied unchanged later) is listed, so it is guarded too.""" + unlisted = sorted( + str(path.relative_to(_TP8)) + for path in _TP8.rglob("*.py") + if "__pycache__" not in path.parts + and str(path.relative_to(_TP8)) not in IDENTICAL_TO_TP4 + and (_TP4 / path.relative_to(_TP8)).is_file() + and not _differs(path, _TP4 / path.relative_to(_TP8)) + ) + assert not unlisted, f"byte copies of the tp4 target not in IDENTICAL_TO_TP4: {unlisted}" 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 04896a8f7a3d..47468e583ab4 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -262,3 +262,21 @@ def test_r1_dep4_with_vanilla_mtp_does_not_match(): modeling_v2_resolve(config, _llm_args("require")) with pytest.raises(ValueError, match="MTP"): modeling_v2_resolve(config, _llm_args("require")) + + +@pytest.mark.parametrize( + "mapping_kwargs, parallel", + [ + (dict(world_size=4, tp_size=4, moe_tp_size=4, moe_ep_size=1), "tp4"), + (dict(world_size=8, tp_size=8, moe_tp_size=8, moe_ep_size=1), "tp8"), + # 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), + ], +) +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.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + assert routing._parallel(Mapping(**mapping_kwargs)) == parallel From bb8674ce7eb364d233b406e051676ca95e6bca43 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:16:44 -0700 Subject: [PATCH 04/14] [None][fix] Drop alignas(64) on MiniMax-M3 CUtensorMap members CUtensorMap already carries its own alignment, which is stricter than 64 bytes in newer CUDA toolkits; clang-hosted nvcc rejects the weaker alignas with "alignment cannot be set to less than the default alignment" (SBSA build). Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/minimaxM3/attnC1.h | 4 ++-- cpp/tensorrt_llm/kernels/minimaxM3/attnMega.h | 4 ++-- cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.h | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/attnC1.h b/cpp/tensorrt_llm/kernels/minimaxM3/attnC1.h index b5f57301ed8e..f13de72e14fb 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3/attnC1.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3/attnC1.h @@ -95,8 +95,8 @@ struct MinimaxM3AttnC1Params long long pf_a0, pf_b0, pf_a1, pf_b1; // mega: L2 prefetch regions issued in the attention window (bytes 0: none) int pf_hint; // 1: evict_last policy on the prefetch uint8_t const* wpre; // mega: optional pre-tiled QKV weight [N/128][K/128][16 KB swizzled smem image], or null - alignas(64) CUtensorMap wmap; // the [2560, 6144] E4M3 weight, 128 x 128 boxes, 128-byte swizzle - alignas(64) CUtensorMap wmap_o; // mega: the [6144, 2048] E4M3 o_proj weight, same boxes + CUtensorMap wmap; // the [2560, 6144] E4M3 weight, 128 x 128 boxes, 128-byte swizzle + CUtensorMap wmap_o; // mega: the [6144, 2048] E4M3 o_proj weight, same boxes }; //! Launch the c1 kernel on `stream`. The caller fills `params` except `wmap_o`, the o_proj weight's TMA diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/attnMega.h b/cpp/tensorrt_llm/kernels/minimaxM3/attnMega.h index d8e267d6e3db..cc751911b3c0 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3/attnMega.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3/attnMega.h @@ -96,8 +96,8 @@ struct MinimaxM3AttnMegaParams unsigned int* flags; // mega: [0, 20) tile counters (4 per call), [32 + CTA] the attention CTAs' targets long long pf_a0, pf_b0, pf_a1, pf_b1; // mega: L2 prefetch regions issued in the attention window (bytes 0: none) int pf_hint; // 1: evict_last policy on the prefetch - alignas(64) CUtensorMap wmap; // the [2560, 6144] E4M3 weight, 128 x 128 boxes, 128-byte swizzle - alignas(64) CUtensorMap wmap_o; // mega: the [6144, 2048] E4M3 o_proj weight, same boxes + CUtensorMap wmap; // the [2560, 6144] E4M3 weight, 128 x 128 boxes, 128-byte swizzle + CUtensorMap wmap_o; // mega: the [6144, 2048] E4M3 o_proj weight, same boxes }; //! Launch the megakernel on `stream`. The caller fills `params` except `wmap`, `wmap_o`, `packed`, diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.h b/cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.h index 663c2d0aad91..e1967d82f556 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3/fusedDecodeAttn.h @@ -86,8 +86,8 @@ struct MinimaxM3FusedDecodeAttnParams long long pf_a0, pf_b0, pf_a1, pf_b1; // mega: L2 prefetch regions issued in the attention window (bytes 0: none) int pf_hint; // 1: evict_last policy on the prefetch uint8_t const* wpre; // mega: optional pre-tiled QKV weight [N/128][K/128][16 KB swizzled smem image], or null - alignas(64) CUtensorMap wmap; // the [2560, 6144] E4M3 weight, 128 x 128 boxes, 128-byte swizzle - alignas(64) CUtensorMap wmap_o; // mega: the [6144, 2048] E4M3 o_proj weight, same boxes + CUtensorMap wmap; // the [2560, 6144] E4M3 weight, 128 x 128 boxes, 128-byte swizzle + CUtensorMap wmap_o; // mega: the [6144, 2048] E4M3 o_proj weight, same boxes }; //! Whether the (decode_query_len, query tokens per cluster, KV blocks per CTA) geometry is compiled. The MXFP8 From a427f563f39a22944cd13323dd4c3fcf8d9488c5 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:33:41 -0700 Subject: [PATCH 05/14] [None][fix] Fix MiniMax-M3 megakernel PR unit tests on current main - Register fake implementations for the new step_inputs_scatter and MiniMax-M3 router, low-M GEMM, producer and allreduce-norm ops. - Give the warmup tests' stand-in engines the default runtime hooks the decoder runner now reads. - Leave the identity table unset in the piecewise-attention metadata mock. - Stub the eager K/V view resolution only while the runtime-policy fixture builds its manager, so the pool-view tests still exercise get_buffers. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../_torch/custom_ops/cpp_custom_ops.py | 54 +++++++++++++++++++ .../sparse/msa/test_msa_runtime_policy.py | 12 +++-- .../test_pytorch_model_engine_warmup.py | 8 +++ .../unittest/_torch/models/test_minimax_m3.py | 1 + 4 files changed, 72 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index f500034c4df7..2521e2faabec 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -398,6 +398,60 @@ def _( return scores.new_empty((scores.shape[2], scores.shape[0], topk), dtype=torch.int32) + @torch.library.register_fake("trtllm::step_inputs_scatter") + def _(stage): + # Writes the staged rows into their destinations; returns nothing. + return None + + @torch.library.register_fake("trtllm::minimax_m3_mxfp8_low_m_gemm") + def _(x, w, wsf, warps, cluster_k, row_tiles, pdl): + return x.new_empty((x.shape[0], w.shape[0])) + + @torch.library.register_fake("trtllm::minimax_m3_router_gemv") + def _(x, w, rows, pdl): + return x.new_empty((x.shape[0], w.shape[0]), dtype=torch.float32) + + def _minimax_m3_router_fp4_outputs(x, w): + m, k = x.shape + return [ + x.new_empty((m, w.shape[0]), dtype=torch.float32), + x.new_empty((m, k // 2), dtype=torch.uint8), + x.new_empty((m * (k // 16), ), dtype=torch.uint8), + ] + + @torch.library.register_fake("trtllm::minimax_m3_router_gemv_fp4") + def _(x, w, sf_scale, pdl): + return _minimax_m3_router_fp4_outputs(x, w) + + @torch.library.register_fake("trtllm::minimax_m3_router_tc") + def _(x, w, sf_scale, pdl): + outputs = _minimax_m3_router_fp4_outputs(x, w) + return outputs if sf_scale is not None else outputs[:1] + + @torch.library.register_fake("trtllm::minimax_m3_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, heads_q, pdl, early): + num_tokens = packed.shape[0] + return [ + packed.new_empty((num_tokens, heads_q, 128), + dtype=torch.float8_e4m3fn), + packed.new_empty((num_tokens, 1, 128), dtype=torch.float8_e4m3fn), + ] + + @torch.library.register_fake("trtllm::minimax_m3_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): + 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 + @torch.library.register_fake("trtllm::kda_decode") def _(x_q: torch.Tensor, x_k: torch.Tensor, x_v: torch.Tensor, w_q_t: torch.Tensor, w_k_t: torch.Tensor, w_v_t: torch.Tensor, 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 01c22219d453..e9bbc98acd07 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 @@ -48,6 +48,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 def no_index_k_buffer(self, layer_idx, **kwargs): del self, layer_idx, kwargs @@ -56,9 +58,13 @@ 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) - return MiniMaxM3KVCacheManagerV2( - num_layers=4, sparse_attention_config=MiniMaxM3SparseAttentionConfig() - ) + # 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) + return MiniMaxM3KVCacheManagerV2( + num_layers=4, sparse_attention_config=MiniMaxM3SparseAttentionConfig() + ) # --------------------------------------------------------------------------- diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py index 83d37b1c22af..708e2e76ac74 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -42,6 +42,7 @@ from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType +from tensorrt_llm._torch.pyexecutor.runtime_hooks import DEFAULT_RUNTIME_HOOKS from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.pyexecutor.warmup_timer import _WarmupTimer from tensorrt_llm._torch.speculative.utils import resolve_draft_len, update_draft_len @@ -374,6 +375,7 @@ def test_compiled_mxfp8_warmup_backend_selection( monkeypatch.setitem(sys.modules, "flashinfer", SimpleNamespace(autotune=flashinfer_tune)) method = MXFP8LinearMethod() engine = SimpleNamespace( + runtime_hooks=DEFAULT_RUNTIME_HOOKS, _warmup_timer=_WarmupTimer(rank=0), _config=SimpleNamespace( enable_autotuner=True, @@ -491,6 +493,7 @@ def test_autotuner_warmup_skips_only_completed_token_shape( ) -> None: """A repeated warmup with an already tuned token shape must not re-run tuning.""" engine = SimpleNamespace( + runtime_hooks=DEFAULT_RUNTIME_HOOKS, _config=SimpleNamespace( enable_autotuner=True, torch_compile_enabled=False, @@ -1281,6 +1284,7 @@ def flashinfer_autotune(): self.assertEqual(method.backend, "trtllm") engine = SimpleNamespace( + runtime_hooks=DEFAULT_RUNTIME_HOOKS, _warmup_timer=_WarmupTimer(rank=0), _config=SimpleNamespace(enable_autotuner=False, torch_compile_enabled=False), cuda_graph_runner=SimpleNamespace(enabled=True), @@ -1344,6 +1348,7 @@ def flashinfer_autotune(): self.assertFalse(method.needs_native_autotune) engine = SimpleNamespace( + runtime_hooks=DEFAULT_RUNTIME_HOOKS, _warmup_timer=_WarmupTimer(rank=0), _config=SimpleNamespace( enable_autotuner=True, @@ -1468,6 +1473,7 @@ def flashinfer_autotune(): os.environ.pop("TLLM_AUTOTUNER_CACHE_PATH", None) method = MXFP8LinearMethod() engine = SimpleNamespace( + runtime_hooks=DEFAULT_RUNTIME_HOOKS, _warmup_timer=_WarmupTimer(rank=0), _config=SimpleNamespace( enable_autotuner=True, @@ -1577,6 +1583,7 @@ def test_flashinfer_mxfp8_rank_mismatch_falls_back_before_warmup(self): os.environ.pop("TRTLLM_MXFP8_GEMM_BACKEND", None) method = MXFP8LinearMethod() engine = SimpleNamespace( + runtime_hooks=DEFAULT_RUNTIME_HOOKS, _warmup_timer=_WarmupTimer(rank=0), _config=SimpleNamespace( enable_autotuner=True, @@ -1655,6 +1662,7 @@ def test_native_mxfp8_respects_disabled_global_autotuner(self): os.environ.pop("TRTLLM_MXFP8_GEMM_BACKEND", None) method = MXFP8LinearMethod() engine = SimpleNamespace( + runtime_hooks=DEFAULT_RUNTIME_HOOKS, _warmup_timer=_WarmupTimer(rank=0), _config=SimpleNamespace(enable_autotuner=False, torch_compile_enabled=False), cuda_graph_runner=SimpleNamespace(enabled=False), diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index 2ebdde977ca5..cef2c1b548d9 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -985,6 +985,7 @@ def test_piecewise_attention_boundary_preserves_indexer_cache_contract( msa_write_idx_k=Mock(), msa_prefill_proxy_plan=None, msa_prefill_n_valid_blocks=None, + _msa_prefill_identity=None, msa_kv_indices=torch.tensor([0, 1], dtype=torch.int32), msa_qo_lens_cpu=torch.tensor([2], dtype=torch.int32), msa_kv_lens_cpu=torch.tensor([2], dtype=torch.int32), From 8cb295c4e927fd68cf8cf6a2e4d70b968e6ed7bc Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Sun, 18 Oct 2026 17:47:16 -0700 Subject: [PATCH 06/14] [None][refactor] MiniMax-M3 modeling_v2: flat TP4/TP8 target layout Restructure the minimax_m3_nvfp4__sm_103__tp4 and __tp8 targets into the flat layout the modeling_v2 exemplars use, with no change to what they compute. Each target owns its own copy of every file (the targets share no code): * modeling.py: the dispatch tables (FAST_PATHS, row ranges, miss/uncovered policy), the flat core, PrefillTarget / DecodeTarget, the one-model Eagle3 drafter and worker, the step fold, runtime hooks and the shell. The core selects the target on num_contexts alone: every pure generation step is DecodeTarget's (its decode kernels up to 48 tokens, its own generic layers for wider steps or embeddings), every step with context rows PrefillTarget's. The two targets share no body: each owns its layer, attention and MoE compositions, and its step_args is the one projection of the MSA metadata its forward reads (DecodeTarget states num_contexts=0 / num_ctx_tokens=0). Shared module-level code is limited to single-call op adapters that read no step metadata. * weights.py: the language-model and drafter weight tables (host-side TP slicing, consumed checkpoint keys released during the load; at TP8 the 4 KV / index heads are duplicated on rank pairs as the built-in loader does). * ops/: catalog entries with certified cells for every call the model makes; kernel wrappers route through their entries. The MSA and drafter attention entries carry a numeric reference and GPU tests that drive every cell against real caches and prepared metadata. * kernels/: the CuTe DSL GEMV, MoE engine and embedding+first-norm kernels move into each target (the old shared cute_dsl_kernels/minimax_m3 copy is removed); capacity annotations on every constant. * The legacy dispatch/layers/linears/runtime/spec_fold/spec_worker modules are removed; the unit tests are ported to the flat targets. * The unfused QKV / index projection runtime (fuse_qkv_index_projection=False, FP8 indexer, FP8 or NVFP4 main KV) runs as the built-in model does. Routing (models/minimax_m3/routing.py) now declines what the targets do not build or run, so modeling_v2="auto" falls back to the built-in model and "require" names the missed criterion: * Eagle3 drafters other than the targets' one-layer drafter (read from the drafter checkpoint's config.json). * Runtimes other than MSA with an FP8 indexer and an FP8 or NVFP4 main KV cache. The resolver hands routing modules that set ROUTES_ON_RUNTIME a RuntimeFacts (sparse attention config and the KV cache quantization the engine will run with); only MiniMax-M3 sets it. _validate_runtime stays as the backstop. Adds a whole-model accuracy test under modeling_v2="require" (GSM8K without speculation, and with one-model Eagle3 plus a chat-GSM8K acceptance probe), on the 4-GPU GB300 list (TP4) and the 2-node GB300 list (TP8). Bitwise identical to the previous targets on GB300 (tokens and per-step outputs on every rank, no speculation and Eagle3, fused and unfused projections, TP4 and TP8, including pure generation steps wider than the decode kernels' 48 tokens). Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../_experimental/modeling_v2/__init__.py | 2 + .../modeling_v2/_router_index.py | 48 +- .../minimax_m3_nvfp4__sm_103__tp4/dispatch.py | 580 -- .../kernels}/embedding_firstnorm/__init__.py | 0 .../m3_embedding_firstnorm_kernel.py | 2 +- .../kernels}/embedding_firstnorm/op.py | 26 +- .../kernels/long_select.py | 2 +- .../kernels/m3_attn_lc.py | 3 +- .../kernels/m3_attn_mega.py | 17 +- .../kernels/m3_attn_mega_c1.py | 19 +- .../kernels/m3_draft_step.py | 13 +- .../kernels/m3_force_vec.py | 35 + .../kernels/m3_fused_attn.py | 13 +- .../kernels/m3_index_score.py | 4 +- .../kernels/m3_mc_allreduce.py | 2 +- .../kernels/m3_moe}/__init__.py | 0 .../kernels/m3_moe}/m3_moe16_kernel.py | 36 +- .../kernels/m3_moe}/m3_moe_kernel.py | 41 +- .../kernels/m3_moe}/op.py | 19 +- .../kernels/m3_moe}/op16.py | 0 .../kernels/m3_router_gemv.py | 2 +- .../kernels/m3_sampled_accept.py | 15 +- .../kernels/m3_shard_accept.py | 10 +- .../kernels/m3_spec_tail.py | 21 +- .../kernels}/mx_gemv/__init__.py | 0 .../kernels}/mx_gemv/m3_mx_gemv_kernel.py | 8 +- .../kernels}/mx_gemv/op.py | 4 +- .../kernels/mxfp8_low_m_gemm.py | 2 +- .../kernels/triton_sparse_decode.py | 10 +- .../minimax_m3_nvfp4__sm_103__tp4/layers.py | 1682 ------ .../minimax_m3_nvfp4__sm_103__tp4/linears.py | 259 - .../minimax_m3_nvfp4__sm_103__tp4/modeling.py | 5086 +++++++++++++++-- .../ops/activation}/__init__.py | 2 +- .../ops/activation/m3_silu_and_mul.py | 83 + .../ops/attention/m3_draft_attention.py | 111 + .../attention/m3_fp8_indexer_qk_norm_rope.py | 114 + .../ops/attention/m3_msa_attention.py | 224 + .../ops/attention/m3_qk_norm_rope_fp8.py | 137 + .../ops/attention/m3_qkv_insert.py | 134 + .../ops/attention/m3_select_blocks.py | 82 + .../ops/comm/m3_allreduce.py | 61 + .../ops/embedding/__init__.py | 3 + .../ops/embedding/m3_embedding.py | 48 + .../ops/gemm/m3_bf16_linear.py | 41 + .../ops/gemm/m3_fp32_linear.py | 46 + .../ops/gemm/m3_mx_gemv.py | 2 +- .../ops/gemm/m3_mx_gemv_bf16.py | 2 +- .../ops/gemm/m3_mxfp8_gemm.py | 58 + .../ops/gemm/m3_mxfp8_linear.py | 137 + .../ops/moe/m3_moe16_engine.py | 4 +- .../ops/moe/m3_moe_engine.py | 11 +- .../ops/norm/m3_embedding_firstnorm.py | 6 +- .../ops/spec/m3_force_vec.py | 86 + .../minimax_m3_nvfp4__sm_103__tp4/runtime.py | 50 - .../spec_fold.py | 112 - .../spec_worker.py | 1166 ---- .../minimax_m3_nvfp4__sm_103__tp4/weights.py | 899 ++- .../minimax_m3_nvfp4__sm_103__tp8/dispatch.py | 527 -- .../kernels/embedding_firstnorm/__init__.py | 3 + .../m3_embedding_firstnorm_kernel.py | 128 + .../kernels/embedding_firstnorm/op.py | 142 + .../kernels/long_select.py | 2 +- .../kernels/m3_attn_lc8.py | 16 +- .../kernels/m3_draft_lowm.py | 121 - .../kernels/m3_draft_step.py | 7 +- .../kernels/m3_force_vec.py | 35 + .../kernels/m3_index_score.py | 4 +- .../kernels/m3_mc_allreduce.py | 2 +- .../kernels/m3_moe/m3_moe_kernel.py | 39 +- .../kernels/m3_moe/op.py | 24 +- .../kernels/m3_router_gemv.py | 2 +- .../kernels/m3_sampled_accept.py | 15 +- .../kernels/m3_shard_accept.py | 10 +- .../kernels/m3_spec_tail.py | 21 +- .../kernels/mx_gemv/m3_mx_gemv_kernel.py | 8 +- .../kernels/mx_gemv/op.py | 4 +- .../kernels/triton_sparse_decode.py | 18 +- .../minimax_m3_nvfp4__sm_103__tp8/layers.py | 1403 ----- .../minimax_m3_nvfp4__sm_103__tp8/linears.py | 232 - .../minimax_m3_nvfp4__sm_103__tp8/modeling.py | 4731 +++++++++++++-- .../ops/activation/__init__.py | 3 + .../ops/activation/m3_silu_and_mul.py | 83 + .../ops/attention/m3_draft_attention.py | 111 + .../attention/m3_fp8_indexer_qk_norm_rope.py | 114 + .../ops/attention/m3_msa_attention.py | 224 + .../ops/attention/m3_qk_norm_rope_fp8.py | 137 + .../ops/attention/m3_qkv_insert.py | 134 + .../ops/attention/m3_select_blocks.py | 82 + .../ops/comm/m3_allreduce.py | 60 + .../ops/embedding/__init__.py | 3 + .../ops/embedding/m3_embedding.py | 48 + .../ops/gemm/m3_bf16_linear.py | 41 + .../ops/gemm/m3_draft_lowm.py | 36 +- .../ops/gemm/m3_fp32_linear.py | 46 + .../ops/gemm/m3_mxfp8_gemm.py | 58 + .../ops/gemm/m3_mxfp8_linear.py | 137 + .../ops/norm/m3_embedding_firstnorm.py | 6 +- .../ops/spec/m3_force_vec.py | 86 + .../minimax_m3_nvfp4__sm_103__tp8/runtime.py | 50 - .../spec_fold.py | 112 - .../spec_worker.py | 1119 ---- .../minimax_m3_nvfp4__sm_103__tp8/weights.py | 917 ++- .../modeling_v2/models/minimax_m3/routing.py | 126 +- .../accuracy/modeling_v2/test_minimax_m3.py | 168 + .../test-db/l0_gb300_multi_gpus.yml | 2 + .../l0_gb300_multi_nodes_node2_gpu8.yml | 3 + .../test_modeling_v2_m3_silu_and_mul.py | 49 + .../attention/_m3_attn_runtime_tp4.py | 274 + .../test_modeling_v2_m3_dense_decode.py | 2 +- .../test_modeling_v2_m3_draft_attention.py | 127 + ...modeling_v2_m3_fp8_indexer_qk_norm_rope.py | 49 + .../test_modeling_v2_m3_msa_attention.py | 214 + .../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 + ...p_matrix.py => _m3_allreduce_op_matrix.py} | 84 +- ...rm.py => test_modeling_v2_m3_allreduce.py} | 14 +- ...eling_v2_m3_mc_allreduce_norm_op_matrix.py | 18 - ...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_fp32_linear.py | 32 + .../gemm/test_modeling_v2_m3_mxfp8_gemm.py | 38 + .../gemm/test_modeling_v2_m3_mxfp8_linear.py | 46 + .../moe/_m3_router.py | 86 + .../moe/test_modeling_v2_m3_router_gemv.py | 134 +- .../test_modeling_v2_m3_router_gemv_fp4.py | 40 + .../moe/test_modeling_v2_m3_router_tc.py | 55 + .../spec/test_modeling_v2_m3_draft_step.py | 13 +- .../test_modeling_v2_m3_exact_numerics.py | 26 +- .../spec/test_modeling_v2_m3_force_vec.py | 47 + .../test_modeling_v2_m3_capture_boundary.py | 126 +- .../test_modeling_v2_m3_dispatch.py | 318 +- .../test_modeling_v2_m3_prebuild.py | 234 +- .../test_modeling_v2_m3_rank_agreement.py | 59 +- .../test_modeling_v2_m3_silu_and_mul.py | 49 + .../attention/_m3_attn_runtime_tp8.py | 274 + .../test_modeling_v2_m3_dense_decode.py | 2 +- .../test_modeling_v2_m3_draft_attention.py | 127 + ...modeling_v2_m3_fp8_indexer_qk_norm_rope.py | 49 + .../test_modeling_v2_m3_fused_score_select.py | 2 +- .../test_modeling_v2_m3_msa_attention.py | 214 + .../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_sparse_decode.py | 8 +- ...p_matrix.py => _m3_allreduce_op_matrix.py} | 84 +- ...rm.py => test_modeling_v2_m3_allreduce.py} | 14 +- ...eling_v2_m3_mc_allreduce_norm_op_matrix.py | 18 - ...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_draft_lowm.py | 82 +- .../gemm/test_modeling_v2_m3_fp32_linear.py | 32 + .../gemm/test_modeling_v2_m3_mxfp8_gemm.py | 38 + .../gemm/test_modeling_v2_m3_mxfp8_linear.py | 46 + .../moe/_m3_router.py | 86 + .../moe/test_modeling_v2_m3_router_gemv.py | 134 +- .../test_modeling_v2_m3_router_gemv_fp4.py | 40 + .../moe/test_modeling_v2_m3_router_tc.py | 55 + .../spec/test_modeling_v2_m3_draft_step.py | 6 +- .../test_modeling_v2_m3_exact_numerics.py | 50 +- .../spec/test_modeling_v2_m3_force_vec.py | 47 + .../test_modeling_v2_m3_capture_boundary.py | 126 +- .../test_modeling_v2_m3_dispatch.py | 394 +- .../test_modeling_v2_m3_prebuild.py | 205 +- .../test_modeling_v2_m3_rank_agreement.py | 63 +- .../test_modeling_v2_m3_tp4_drift.py | 46 +- .../modeling_v2/test_modeling_v2_routing.py | 278 + 171 files changed, 18411 insertions(+), 9927 deletions(-) delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/dispatch.py rename tensorrt_llm/_torch/{cute_dsl_kernels/minimax_m3 => _experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels}/embedding_firstnorm/__init__.py (100%) rename tensorrt_llm/_torch/{cute_dsl_kernels/minimax_m3 => _experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels}/embedding_firstnorm/m3_embedding_firstnorm_kernel.py (98%) rename tensorrt_llm/_torch/{cute_dsl_kernels/minimax_m3 => _experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels}/embedding_firstnorm/op.py (86%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_force_vec.py rename tensorrt_llm/_torch/{cute_dsl_kernels/minimax_m3/moe => _experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe}/__init__.py (100%) rename tensorrt_llm/_torch/{cute_dsl_kernels/minimax_m3/moe => _experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe}/m3_moe16_kernel.py (99%) rename tensorrt_llm/_torch/{cute_dsl_kernels/minimax_m3/moe => _experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe}/m3_moe_kernel.py (99%) rename tensorrt_llm/_torch/{cute_dsl_kernels/minimax_m3/moe => _experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe}/op.py (97%) rename tensorrt_llm/_torch/{cute_dsl_kernels/minimax_m3/moe => _experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe}/op16.py (100%) rename tensorrt_llm/_torch/{cute_dsl_kernels/minimax_m3 => _experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels}/mx_gemv/__init__.py (100%) rename tensorrt_llm/_torch/{cute_dsl_kernels/minimax_m3 => _experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels}/mx_gemv/m3_mx_gemv_kernel.py (99%) rename tensorrt_llm/_torch/{cute_dsl_kernels/minimax_m3 => _experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels}/mx_gemv/op.py (99%) delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/layers.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/linears.py rename tensorrt_llm/_torch/{cute_dsl_kernels/minimax_m3 => _experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/activation}/__init__.py (60%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/activation/m3_silu_and_mul.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_draft_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp4/ops/attention/m3_msa_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp4/ops/attention/m3_qkv_insert.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_select_blocks.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_allreduce.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding/m3_embedding.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_bf16_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_fp32_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_gemm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_force_vec.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/runtime.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec_fold.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec_worker.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/dispatch.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/embedding_firstnorm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/embedding_firstnorm/op.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_draft_lowm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_force_vec.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/layers.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/linears.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/activation/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/activation/m3_silu_and_mul.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_draft_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp8/ops/attention/m3_msa_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp8/ops/attention/m3_qkv_insert.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_select_blocks.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_allreduce.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding/m3_embedding.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_bf16_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_fp32_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mxfp8_gemm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mxfp8_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_force_vec.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/runtime.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec_fold.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec_worker.py create mode 100644 tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp4/attention/_m3_attn_runtime_tp4.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_draft_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp4/attention/test_modeling_v2_m3_msa_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp4/attention/test_modeling_v2_m3_qkv_insert.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_select_blocks.py rename tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/{_m3_draft_ar_norm_op_matrix.py => _m3_allreduce_op_matrix.py} (78%) rename tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/{test_modeling_v2_m3_draft_ar_norm.py => test_modeling_v2_m3_allreduce.py} (55%) create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp4/embedding/test_modeling_v2_m3_embedding.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/_m3_mxfp8_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_bf16_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_fp32_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/_m3_router.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_gemv_fp4.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_tc.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_force_vec.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp8/attention/_m3_attn_runtime_tp8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_draft_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp8/attention/test_modeling_v2_m3_msa_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp8/attention/test_modeling_v2_m3_qkv_insert.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_select_blocks.py rename tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/{_m3_draft_ar_norm_op_matrix.py => _m3_allreduce_op_matrix.py} (79%) rename tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/{test_modeling_v2_m3_draft_ar_norm.py => test_modeling_v2_m3_allreduce.py} (66%) create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/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_103__tp8/embedding/test_modeling_v2_m3_embedding.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/_m3_mxfp8_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_bf16_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_fp32_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/_m3_router.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_gemv_fp4.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_tc.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_force_vec.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py index 7dc4b711fd91..a9c6c10e7552 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py @@ -31,6 +31,7 @@ MODELING_V2_VALIDATE_ENV, ModelingV2Mode, NoTarget, + RuntimeFacts, modeling_v2_resolve, resolved_mode, step_contract_enabled, @@ -41,6 +42,7 @@ "MODELING_V2_VALIDATE_ENV", "ModelingV2Mode", "NoTarget", + "RuntimeFacts", "modeling_v2_resolve", "resolved_mode", "step_contract_enabled", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index 6dc4a53867bb..e8c3824fdeeb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -18,7 +18,9 @@ * The routing module's ``route`` owns the rest of the decision as one forward-reading tree over the facts that fix which code runs: the checkpoint's config, the parallel mapping, the GPU's compute capability and - the speculative-decoding config. It returns the target class, or raises + the speculative-decoding config (and, for a module that sets + ``ROUTES_ON_RUNTIME``, the engine arguments in :class:`RuntimeFacts`). It + returns the target class, or raises :class:`NoTarget` with the reason in words. Reading that one file tells you where any configuration lands. @@ -36,8 +38,9 @@ import enum import os +from dataclasses import dataclass from importlib import import_module -from typing import TYPE_CHECKING, Dict, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple from tensorrt_llm.logger import logger @@ -105,6 +108,36 @@ class NoTarget(Exception): """ +@dataclass(frozen=True) +class RuntimeFacts: + """Engine arguments that fix which code runs but are not in the checkpoint. + + A routing module that sets ``ROUTES_ON_RUNTIME = True`` receives them as + ``route(..., runtime=RuntimeFacts)`` and declines (``NoTarget``) a runtime + its targets are not written for, so ``auto`` falls back to the built-in + implementation instead of the target refusing at construction. + """ + + # ``ModelConfig.sparse_attention_config`` (the LLM API's sparse_attention_config), or None. + sparse_attention_config: Any + # The main KV cache quantization the engine will run with (a ``QuantAlgo`` value, None for the + # unquantized cache): ``kv_cache_config.dtype``, or the checkpoint's under ``auto``. + kv_cache_quant_algo: Optional[str] + + @classmethod + def of(cls, config: "ModelConfig", llm_args: "TorchLlmArgs") -> "RuntimeFacts": + # The resolver runs before the loader applies kv_cache_config.dtype to the quant config + # (validate_and_set_kv_cache_quant); this is the value that function will set. + from tensorrt_llm._torch.pyexecutor.model_loader import _KV_CACHE_MAP + + dtype = llm_args.kv_cache_config.dtype + checkpoint = config.quant_config.kv_cache_quant_algo if config.quant_config else None + return cls( + sparse_attention_config=config.sparse_attention_config, + kv_cache_quant_algo=checkpoint if dtype == "auto" else _KV_CACHE_MAP.get(dtype), + ) + + # The mode each target class was last resolved under (the loader builds the class right after resolving it). _RESOLVED_MODES: Dict[type, ModelingV2Mode] = {} @@ -161,8 +194,17 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti raise NoTarget( f"no routing module for this architecture; routed architectures: {known}" ) + runtime = ( + {"runtime": RuntimeFacts.of(config, llm_args)} + if getattr(routing, "ROUTES_ON_RUNTIME", False) + else {} + ) target = routing.route( - pretrained_config, config.mapping, _device_compute_capability(), config.spec_config + pretrained_config, + config.mapping, + _device_compute_capability(), + config.spec_config, + **runtime, ) except NoTarget as why: if mode is ModelingV2Mode.REQUIRE: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/dispatch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/dispatch.py deleted file mode 100644 index 47b00f57275d..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/dispatch.py +++ /dev/null @@ -1,580 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The decode-step fast paths of this target: the rows each one takes, the catalog entry that certifies them, and -what happens when one cannot run. - -One target serves every concurrency with the same code; which fast path a step takes is a pure function of its -shape, written here once (``plan_for``, ``moe_path``, ``router_path``, ``attention_path``, ``linear_path``) and -called by the forward. Every row range is a named constant below, and ``FAST_PATHS`` pairs each with the catalog -entry whose ``CELLS`` certify it: ``test_modeling_v2_m3_dispatch.py`` checks that every row count (or shape key) a -range covers is one of the entry's cells. - -Two kinds of not taking a fast path, kept apart: - -* **uncovered**: the step is outside every covered range -- prefill, mixed steps, generation wider than the ranges, - an eager decode step (the GEMV dispatch is a decode-graph one), a long-context request, an Eagle3 capture layer - that runs the generic layer by design. ``uncovered`` logs the reason once per (path, reason); the caller runs the - generic path in every mode. -* **missed**: the step is inside a covered range but a precondition is missing -- a workspace, a push descriptor or - a compiled engine not built, or a kernel declining operands it is certified for. Under ``modeling_v2='require'`` - ``missed`` raises ``FastPathUnavailable`` naming the path, the rows and the precondition; otherwise it logs once - per (path, reason) and the caller falls back. - -Compiled kernels (the MoE engines, the decode GEMVs with their push variants, the index scorer) are built by -``build_targets`` 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, kernel modules) 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``) rather -than a silent capture of the fallback. -""" - -from __future__ import annotations - -import threading -from dataclasses import dataclass -from types import MappingProxyType -from typing import Any, Callable, Dict, FrozenSet, Hashable, Iterable, Mapping, Optional, Tuple - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper -from tensorrt_llm.logger import logger - -from .ops.attention.m3_attn_c1 import m3_attn_c1 -from .ops.attention.m3_attn_mega import m3_attn_mega -from .ops.attention.m3_dense_decode import m3_dense_decode -from .ops.attention.m3_fused_decode_attn import m3_fused_decode_attn -from .ops.attention.m3_index_score import m3_index_score -from .ops.attention.m3_producer import m3_producer -from .ops.attention.m3_select_long import MAX_BLOCKS as SELECT_LONG_MAX_BLOCKS -from .ops.attention.m3_select_long import m3_select_long -from .ops.attention.m3_sparse_decode import m3_sparse_decode -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.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_target_head import MAX_ROWS as TARGET_HEAD_KERNEL_MAX_ROWS -from .ops.gemm.m3_target_head import m3_target_head -from .ops.gemm.mxfp8_low_m_gemm import mxfp8_low_m_gemm -from .ops.moe.m3_moe16_engine import m3_moe16_engine -from .ops.moe.m3_moe_engine import m3_moe_engine -from .ops.moe.m3_router_gemv import m3_router_gemv -from .ops.moe.m3_router_gemv_fp4 import m3_router_gemv_fp4 -from .ops.moe.m3_router_tc import m3_router_tc -from .ops.norm.m3_embedding_firstnorm import m3_embedding_firstnorm -from .ops.spec.m3_draft_input import m3_draft_input -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_shard_accept import m3_shard_accept - -# -------------------------------------------------------------------------------------------------------------------- -# Row ranges -# -------------------------------------------------------------------------------------------------------------------- - -# The largest generation step the decode target takes: 8 requests verifying 1 + 3 drafts and their padding rows. -DECODE_MAX_TOKENS = 48 - - -@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 (kernels/m3_mc_allreduce.BUCKETS). -DECODE_PLANS = ( - DecodePlan("tiny", 8, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), - DecodePlan("small", 16, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), - DecodePlan("verify", DECODE_MAX_TOKENS, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), -) - -# 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. -DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({1: 16384}) - -# 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 - -# The fused embedding + first RMSNorm kernel: one owner per row count and index dtype. -EMBED_FIRSTNORM_MAX_TOKENS = 8 -# 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 = 16 -AR_PUSH_MODE = 2 -# The 1..8-token MoE engine (router, routed experts and the shared expert in one launch) pushes its sum into the MoE -# boundary's all-reduce (prepushed mode 2). -MOE_ENGINE_MAX_TOKENS = 8 -MOE_PUSH_MAX_TOKENS = MOE_ENGINE_MAX_TOKENS -# The 9..13-row engine (router, routed and shared experts in one kernel); its bf16 sum goes to the boundary. -MOE16_MIN_TOKENS = 9 -MOE16_MAX_TOKENS = 13 -# 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). -MC_FINALIZE_MIN_TOKENS = 9 -# The router: the fp32 GEMV at one row, the tensor-core router (fused with the routed experts' NVFP4 activation -# quantization) at 2..48 rows. GB300 tp4, behind a PDL producer, with the NVFP4 activation: 2.90 / 2.93 / 3.20 / -# 3.82 / 7.36 us at 1 / 4 / 8 / 12 / 48 tokens against the fp32 GEMV's 2.85 / 3.96 / 4.88 / 6.01 / - us. -ROUTER_GEMV_MAX_TOKENS = 1 -ROUTER_TC_MAX_TOKENS = 48 -# The fused norm/RoPE/KV-insert/sparse-decode/merge kernel: up to 8 query rows per request and 16 query rows per -# step (above, its 16-CTA clusters run in waves and the three-kernel path is faster: GB300 -13% fAL spec c8). -FUSED_ATTN_MAX_QUERY_LEN = 8 -FUSED_ATTN_MAX_TOKENS = 16 -# One request verifying a draft of 6 query rows runs attention + o_proj as one launch (96 CTAs: 6 clusters of 16). -ATTN_C1_QUERY_LEN = 6 -# 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 (kernels/m3_shard_accept), at the configured draft length. -SHARD_ACCEPT_DRAFT_LEN = 3 -SHARD_ACCEPT_MAX_REQUESTS = DECODE_MAX_TOKENS // (SHARD_ACCEPT_DRAFT_LEN + 1) -# 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 (kernels/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. -TARGET_HEAD_MAX_ROWS = TARGET_HEAD_KERNEL_MAX_ROWS -# 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 = 8 -# 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 (kernels/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 (kernels/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 = True -DRAFT_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN -DRAFT_STEP_MAX_REQUESTS = DRAFT_STEP_KERNEL_MAX_TOKENS -# The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows -# (kernels/triton_sparse_decode.MQ_MIN_TOKENS). -SPARSE_DECODE_MQ_MIN_TOKENS = 17 -# 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 -LONG_SELECT_MAX_BLOCKS = SELECT_LONG_MAX_BLOCKS -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; on wider steps the early-launched attention CTAs crowd -# the select's CTAs off the SMs (8 requests x 4 rows: +3.1 us per layer), so they keep the trigger at the select's exit. -SELECT_EARLY_MAX_ROWS = 8 -# 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 (larger batches -# are bandwidth bound, where the extra early tiles only crowd the prefetch). -INDEX_SCORE_EARLY_MAX_REQUESTS = 1 - -# 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. -_ALL_ROWS = frozenset(range(1, DECODE_MAX_TOKENS + 1)) -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, -} -# 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, -} -# The fused low-M MXFP8 GEMM (projections marked for it), per projection: its built configurations' row limit -# (kernels/mxfp8_low_m_gemm). The bf16 GEMV serves every one of these rows first; the GEMM takes them only if it -# declines. -LOW_M_GEMM_MAX_TOKENS: Dict[Tuple[int, int], int] = { - QKV_SPARSE: 8, - O_PROJ: 8, - DENSE_GATE_UP: 1, - DENSE_DOWN: 8, -} - -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 _fused_attn_keys() -> FrozenSet[Tuple[int, int]]: - return frozenset( - (qlen, reqs) - for qlen in range(1, FUSED_ATTN_MAX_QUERY_LEN + 1) - for reqs in range(1, FUSED_ATTN_MAX_TOKENS // qlen + 1) - ) - - -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"), -) -# One request decoding one token on an MXFP8 layer input: QKV GEMV + attention + o_proj in one launch. -ATTN_MEGA = RowRange("attn_mega", frozenset({(1, 1)}), m3_attn_mega, _const((1, 1))) -ATTN_C1 = RowRange( - "attn_c1", frozenset({(ATTN_C1_QUERY_LEN, 1)}), m3_attn_c1, _const((ATTN_C1_QUERY_LEN, 1)) -) -FUSED_DECODE_ATTN = RowRange( - "fused_decode_attn", - _fused_attn_keys(), - m3_fused_decode_attn, - lambda spec: ((spec["qlen"], len(spec["seq_lens"])),), -) -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",), -) -LONG_SELECT = RowRange( - "long_select", frozenset(LONG_SELECT_KEY_WIDTHS), m3_select_long, _long_select_widths -) -MC_ALLREDUCE = RowRange( - "mc_allreduce_norm", - _rows(1, DECODE_MAX_TOKENS), - 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), -) -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, DECODE_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"), -) -ROUTER_GEMV = RowRange( - "m3_router_gemv", _rows(1, ROUTER_GEMV_MAX_TOKENS), m3_router_gemv, _spec_rows("m") -) -ROUTER_GEMV_FP4 = RowRange( - "m3_router_gemv_fp4", _rows(1, ROUTER_GEMV_MAX_TOKENS), m3_router_gemv_fp4, _spec_rows("m") -) -ROUTER_TC = RowRange( - "m3_router_tc", - _rows(ROUTER_GEMV_MAX_TOKENS + 1, ROUTER_TC_MAX_TOKENS), - m3_router_tc, - _spec_rows("m"), -) -# 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"]),), -) -LOW_M_GEMM = RowRange( - "mxfp8_low_m_gemm", - frozenset((nk, m) for nk, hi in LOW_M_GEMM_MAX_TOKENS.items() for m in _rows(1, hi)), - mxfp8_low_m_gemm, - lambda spec: (((spec["n"], spec["k"]), 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 (), -) - -# 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"]),), -) -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, - ATTN_MEGA, - ATTN_C1, - FUSED_DECODE_ATTN, - DENSE_DECODE_SPLIT, - SPARSE_DECODE_PRODUCER, - PRODUCER, - INDEX_SCORE, - LONG_SELECT, - MC_ALLREDUCE, - ATTN_PUSH, - MOE_PUSH, - MC_FINALIZE, - MOE_ENGINE, - MOE16_ENGINE, - ROUTER_GEMV, - ROUTER_GEMV_FP4, - ROUTER_TC, - MX_GEMV, - MX_GEMV_BF16, - LOW_M_GEMM, - SHARD_ACCEPT, - TARGET_HEAD, - DRAFT_STEP, - 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 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: 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 - - -def router_path(num_tokens: int) -> Optional[str]: - """The router kernel when the MoE runs unfused: the fp32 GEMV, the tensor-core router, or None (built-in gate).""" - if ROUTER_GEMV.covers(num_tokens): - return ROUTER_GEMV.path - if ROUTER_TC.covers(num_tokens): - return ROUTER_TC.path - return None - - -MX_GEMV_PATH = MX_GEMV.path -MX_GEMV_BF16_PATH = MX_GEMV_BF16.path -LOW_M_GEMM_PATH = LOW_M_GEMM.path -# 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``: the projection is marked for the - fused low-M GEMM. - """ - 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 low_m_gemm and num_tokens <= LOW_M_GEMM_MAX_TOKENS.get((n, k), 0): - return LOW_M_GEMM_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 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() diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/embedding_firstnorm/__init__.py similarity index 100% rename from tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/embedding_firstnorm/__init__.py diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py similarity index 98% rename from tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/m3_embedding_firstnorm_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py index acf27a95316b..78563b4db40d 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/m3_embedding_firstnorm_kernel.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py @@ -13,7 +13,7 @@ NORM_THREADS = 128 NORM_VEC = 8 -NORM_WARPS = 4 +NORM_WARPS = 4 # capacity: fixed: the CTA's warps @cute.kernel diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/embedding_firstnorm/op.py similarity index 86% rename from tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/embedding_firstnorm/op.py index 0acfdecb196e..276f8c32023b 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/embedding_firstnorm/op.py @@ -17,8 +17,8 @@ from .m3_embedding_firstnorm_kernel import m3_embed_norm as _m3_embed_norm_jit -HIDDEN = 6144 -MAX_ROWS = 8 +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] = {} @@ -114,33 +114,29 @@ def m3_embed_norm( 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.""" + 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, norm, rows, id_dtype, use_pdl=True): + 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 ( - norm.weight.dtype == torch.bfloat16 - and norm.weight.shape == (HIDDEN,) - and not norm.use_gemma - ) + 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.norm, self.rows, self.id_dtype, self.use_pdl = ( + self.table, self.weight, self.rows, self.id_dtype, self.use_pdl = ( table, - norm, + weight, rows, id_dtype, use_pdl, ) - self.eps = float(norm.variance_epsilon) + 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, norm.weight, self.eps, self.raw, self.output, use_pdl) + _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.norm.weight, self.eps, self.raw, self.output, self.use_pdl + 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_103__tp4/kernels/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/long_select.py index 4fb5064fa981..c01a82f0da41 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/long_select.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/long_select.py @@ -35,7 +35,7 @@ def long_select_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 (``dispatch.SELECT_EARLY_MAX_ROWS``). + grid wait (``SELECT_EARLY_MAX_ROWS`` in modeling.py). """ if ( max_score_kv.dtype != torch.float32 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_lc.py index 09b4fc4ee57d..88e64f35ceb4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_lc.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_lc.py @@ -16,7 +16,8 @@ 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 (``dispatch.SELECT_EARLY_MAX_ROWS``); +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. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega.py index da6a0ddffcca..964a138afd67 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega.py @@ -52,11 +52,13 @@ import torch +from ..ops.attention.m3_attn_mega import m3_attn_mega as _m3_attn_mega_op + NUM_HEADS = 16 HEAD_DIM = 128 PAGE_SIZE = 128 -MAX_BLOCKS = 16 -PACKED_WIDTH = (NUM_HEADS + 4) * HEAD_DIM +MAX_BLOCKS = 16 # capacity: fixed: the MSA top-16 blocks +PACKED_WIDTH = (NUM_HEADS + 4) * HEAD_DIM # capacity: fixed: the tp4 packed QKV width _module = None @@ -91,8 +93,12 @@ def __init__(self, device, max_tokens: int = 8): ).view(torch.bfloat16) self.flags = torch.zeros(256, dtype=torch.int32, device=device) # Two call-parity buffers of the attention's MXFP8 hand-off to the o_proj units, all 0xFF ("not yet"). - self.q8 = torch.full((2 * max_tokens * 2048,), 255, dtype=torch.uint8, device=device) - self.sf = torch.full((2 * 16 * 512,), 255, dtype=torch.uint8, device=device) + self.q8 = torch.full( + (2 * max_tokens * 2048,), 255, dtype=torch.uint8, device=device + ) # capacity: fixed: a workspace of the kernel's fixed geometry + self.sf = torch.full( + (2 * 16 * 512,), 255, dtype=torch.uint8, device=device + ) # capacity: fixed: a workspace of the kernel's fixed geometry self.calls = 0 # host mirror of the call count (eager use only; graph replays advance the device copy) @@ -156,7 +162,8 @@ def m3_attn_mega( "y must be bf16 [T, 6144]", ) qg, bpc = 1, 1 # one 16-CTA attention cluster, one KV block per CTA - _load_module().minimax_m3_attn_mega( + _load_module() + _m3_attn_mega_op( xq.view(torch.uint8), xsf.view(torch.uint8), qkv_weight.view(torch.uint8), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega_c1.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega_c1.py index e4898442f227..82c3b1cc3593 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega_c1.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega_c1.py @@ -46,11 +46,13 @@ import torch +from ..ops.attention.m3_attn_c1 import m3_attn_c1 as _m3_attn_c1_op + NUM_HEADS = 16 -HEAD_DIM = 128 +HEAD_DIM = 128 # capacity: fixed: the checkpoint's head dim PAGE_SIZE = 128 -MAX_BLOCKS = 16 -PACKED_WIDTH = (NUM_HEADS + 4) * HEAD_DIM +MAX_BLOCKS = 16 # capacity: fixed: the MSA top-16 blocks +PACKED_WIDTH = (NUM_HEADS + 4) * HEAD_DIM # capacity: fixed: the tp4 packed QKV width # The geometries of m3_fused_attn, which default_geometry() reads (only C1_GEOMETRIES is compiled into this op). GEOMETRIES = frozenset( @@ -231,8 +233,12 @@ def __init__(self, device, max_tokens: int = 8): ).view(torch.bfloat16) self.flags = torch.zeros(256, dtype=torch.int32, device=device) # Two call-parity buffers of the attention's MXFP8 hand-off to the o_proj units, all 0xFF ("not yet"). - self.q8 = torch.full((2 * max_tokens * 2048,), 255, dtype=torch.uint8, device=device) - self.sf = torch.full((2 * 16 * 512,), 255, dtype=torch.uint8, device=device) + self.q8 = torch.full( + (2 * max_tokens * 2048,), 255, dtype=torch.uint8, device=device + ) # capacity: fixed: a workspace of the kernel's fixed geometry + self.sf = torch.full( + (2 * 16 * 512,), 255, dtype=torch.uint8, device=device + ) # capacity: fixed: a workspace of the kernel's fixed geometry self.calls = 0 # host mirror of the call count (eager use only; graph replays advance the device copy) @@ -321,7 +327,8 @@ def m3_attn_c1( "y must be bf16 [T, 6144]", ) out2d = output.view(num_tokens, NUM_HEADS * HEAD_DIM) - _load_module().minimax_m3_attn_c1( + _load_module() + _m3_attn_c1_op( packed, ws.flags, kv_data, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_draft_step.py index f15d8dbeff64..31768ebcd2ee 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_draft_step.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_draft_step.py @@ -27,7 +27,6 @@ import torch -from .. import dispatch 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 @@ -63,16 +62,16 @@ def head_supported(lm_head) -> bool: return (n // 16 + _num_sms() - 1) // _num_sms() <= MAX_TILES_PER_CTA -def prepare_head(lm_head, tp_rank: int) -> DraftHead: +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 - ``dispatch.DRAFT_HEAD_CODED``, K is a multiple of 64 and the coding decodes to the shard bit for bit, else (fail + ``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 dispatch.DRAFT_HEAD_CODED and w.shape[1] % 64 == 0: + 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( @@ -104,13 +103,13 @@ def _num_sms() -> int: 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.""" + """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(): - dispatch.capture_missed(dispatch.DRAFT_STEP.path, None, "the draft-step exchange workspace") + # 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 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_force_vec.py new file mode 100644 index 000000000000..b4de37794ca9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_force_vec.py @@ -0,0 +1,35 @@ +# 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``), drawn exactly as its ``_apply_force_accepted_tokens`` does, in one launch. + +Triton: 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. Reached through the ``ops/spec`` entry ``m3_force_vec``. +""" + +import triton +import triton.language as tl + + +@triton.jit +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) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_fused_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_fused_attn.py index c404166fb320..fefd30acf333 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_fused_attn.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_fused_attn.py @@ -42,11 +42,13 @@ import torch -NUM_HEADS = 16 -HEAD_DIM = 128 +from ..ops.attention.m3_fused_decode_attn import m3_fused_decode_attn as _m3_fused_decode_attn_op + +NUM_HEADS = 16 # capacity: fixed: the tp4 query heads +HEAD_DIM = 128 # capacity: fixed: the checkpoint's head dim PAGE_SIZE = 128 -MAX_BLOCKS = 16 -PACKED_WIDTH = (NUM_HEADS + 4) * HEAD_DIM +MAX_BLOCKS = 16 # capacity: fixed: the MSA top-16 blocks +PACKED_WIDTH = (NUM_HEADS + 4) * HEAD_DIM # capacity: fixed: the tp4 packed QKV width # (decode_query_len, tokens per cluster, KV blocks per CTA) geometries compiled into the op: what # default_geometry() returns for every step shape the target routes here (at most 16 tokens). @@ -318,7 +320,8 @@ def m3_fused_decode_attn( (decode_query_len, qg, bpc) in GEOMETRIES, f"geometry {(decode_query_len, qg, bpc)} is not compiled", ) - _load_module().minimax_m3_fused_decode_attn( + _load_module() + _m3_fused_decode_attn_op( packed, kv_data, kv_scale, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_index_score.py index 9748bd1fe4a5..dfe1909e537a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_index_score.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_index_score.py @@ -543,10 +543,10 @@ class IndexScoreRunner: per CTA.""" kernel_cache: dict = {} - SUPPORTED_HEAD_DIM = 128 + 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 + MAX_BLOCK_Q = 32 # capacity: fixed: the scorer's query tile EARLY_PREFETCH_DEPTH = 1 @classmethod diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_mc_allreduce.py index bdb207b7a8ef..cd1ed18a6cb9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_mc_allreduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_mc_allreduce.py @@ -38,7 +38,7 @@ m3_mc_finalize_allreduce_norm as _finalize_allreduce_norm_op, ) -MAX_TOKENS = 48 +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 diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/__init__.py similarity index 100% rename from tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/__init__.py diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe16_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/m3_moe16_kernel.py similarity index 99% rename from tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe16_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/m3_moe16_kernel.py index e317e5e3087a..ff609e78ef29 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe16_kernel.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/m3_moe16_kernel.py @@ -47,11 +47,11 @@ def _cfg(key: str, default): return type(default)(_CFG.get(key, default)) -H = 6144 +H = 6144 # capacity: fixed: the checkpoint's hidden size I_TP = 768 TWO_I = 2 * I_TP E_LOCAL = 128 -TOP_K = 4 +TOP_K = 4 # capacity: fixed: the checkpoint's top-k ROUTE_SCALE = 2.0 SUM_EPS = 1e-20 M_MAX = _cfg("m_max", 12) @@ -66,7 +66,7 @@ def _cfg(key: str, default): LAG = _cfg("lag", 3) G_CAP = TOP_K * M_MAX NG_CAP = M_MAX -THREADS = 512 +THREADS = 512 # capacity: fixed: the CTA size NWARPS = THREADS // 32 GEMV_THREADS = 256 SHARED = bool(_cfg("shared", 0)) @@ -95,12 +95,12 @@ def _cfg(key: str, default): ) # 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_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 +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) +SH_STAGES = SH_GT + 2 * (SH_I // 128) # capacity: fixed: pipeline stages SHX_BAR = 6 SH_PRETILE = False SH_STAGE_BYTES = 16384 + 512 @@ -117,8 +117,8 @@ def _cfg(key: str, default): C1 = 4 U1 = TWO_I // R1 K1_TILES = H // 256 -S1 = K1_TILES // C1 -NUM_BYTES_A = MMA_M * KT_BYTES +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 @@ -150,7 +150,7 @@ def _cfg(key: str, default): 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 +PARTIAL_EMPTY = -1 # capacity: fixed: a sentinel, not a size EPI_THREADS = 128 EPI_BAR = 1 ROUTE_BAR = 2 @@ -161,10 +161,10 @@ def _cfg(key: str, default): 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 +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 @@ -208,7 +208,7 @@ def _cfg(key: str, default): 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 +SHN = 16 # capacity: fixed: a kernel tile constant SH_NPASS = 1 SH_HINT = EVICT_FIRST SH_PF_POST = bool( @@ -238,14 +238,14 @@ def _cfg(key: str, default): 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 +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 +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) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/m3_moe_kernel.py similarity index 99% rename from tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/m3_moe_kernel.py index ac9892f31ed2..1ad72e733c64 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe_kernel.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/m3_moe_kernel.py @@ -62,11 +62,11 @@ def _cfg(key: str, default): # ============================================================================= # Problem shape (trace-time constants). # ============================================================================= -H = 6144 +H = 6144 # capacity: fixed: the checkpoint's hidden size I_TP = 768 TWO_I = 2 * I_TP E_LOCAL = 128 -TOP_K = 4 +TOP_K = 4 # capacity: fixed: the checkpoint's top-k ROUTE_SCALE = 2.0 SUM_EPS = 1e-20 M_MAX = _cfg("m_max", 1) @@ -87,7 +87,7 @@ def _cfg(key: str, default): 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 +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 @@ -108,12 +108,16 @@ def _cfg(key: str, default): 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_KB = H // 32 # gate_up K blocks # capacity: fixed: a kernel tile constant +SH_DKB = SH_I // 32 # down K blocks # capacity: fixed: a kernel tile constant 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_GT = ( + SH_KB // 4 // SH_SPLIT +) # gate_up k-tiles (128 K) per slice # capacity: fixed: a kernel tile constant SH_STAGES_GU = SH_GT -SH_STAGES = SH_GT + 2 * (SH_I // 128) # + 2 down tiles x 6 k-tiles +SH_STAGES = SH_GT + 2 * ( + SH_I // 128 +) # + 2 down tiles x 6 k-tiles # capacity: fixed: pipeline stages 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 @@ -133,14 +137,14 @@ def _cfg(key: str, default): 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 +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 +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 @@ -155,7 +159,7 @@ def _cfg(key: str, default): 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 +PARTIAL_EMPTY = -1 # capacity: fixed: a sentinel, not a size EPI_THREADS = 128 EPI_BAR = 1 ROUTE_BAR = 2 @@ -198,11 +202,12 @@ def _cfg(key: str, default): 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) +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)) @@ -275,8 +280,8 @@ def _cfg(key: str, default): 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 +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 FUSED_ROUTER, "shared expert: fused router only" assert SH_STAGES - SH_NST <= 64, SH_STAGES diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/op.py similarity index 97% rename from tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/op.py index 12f6008d439f..f521a23589ea 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/op.py @@ -52,19 +52,22 @@ import torch -HIDDEN = 6144 +HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size 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) +# The shared expert's intermediate per TP rank. +SHARED_INTER = 768 # capacity: fixed: the shared expert's width per tp4 rank +# m3_moe's shared-expert CTAs (beside num_ctas routed CTAs). +SHARED_CTAS = 24 # capacity: fixed: an SM split, not a size 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) +# 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 -_ITEM_BYTES = _UNIT_BYTES // _S1 +_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). @@ -180,7 +183,9 @@ def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast 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) + out = torch.empty( + n_cta, 36, 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 = torch.where((q & 1) == 0, 32 * s + ri, I_ + 32 * s + ri) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/op16.py similarity index 100% rename from tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op16.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/op16.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_router_gemv.py index 88a475f509d5..30855eb7f05d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_router_gemv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_router_gemv.py @@ -20,7 +20,7 @@ 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 +_MAX_DYNAMIC_SMEM = 227 * 1024 - 1024 # capacity: fixed: the device's shared-memory limit def m3_router_gemv( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_sampled_accept.py index 315e80947213..59c58336c5d7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_sampled_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_sampled_accept.py @@ -28,11 +28,12 @@ import torch -from .. import dispatch +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/minimaxM3/sampledAccept.cu) and called as -# torch.ops.trtllm.minimax_m3_sampled_accept. -MAX_ROWS = 48 +# The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.cu) and reached through +# the ops/spec entry m3_sampled_accept. +MAX_ROWS = 48 # capacity: fixed: the kernel's built rows (decode steps) class SampledAcceptWorkspace: @@ -49,7 +50,7 @@ def __init__(self, uc, mc_ptr, rank, world, handle=None, comm=None): @classmethod def create(cls, mapping, device) -> "SampledAcceptWorkspace": - words = int(torch.ops.trtllm.minimax_m3_sampled_accept_ws_words(max(mapping.tp_size, 1))) + words = int(_ws_words_op(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) @@ -88,7 +89,7 @@ def workspace(mapping, device) -> Optional[SampledAcceptWorkspace]: global _WS if _WS is None: if torch.cuda.is_current_stream_capturing(): - dispatch.capture_missed("spec_tail", None, "the sampled-accept workspace") + # 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) return _WS @@ -131,7 +132,7 @@ def sampled_accept( 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_sampled_accept( + _sampled_accept_op( logits, int(vocab_offset), temperatures[:rows].float().contiguous(), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_shard_accept.py index e1dafd292d6d..2b7d05ad920f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_shard_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_shard_accept.py @@ -13,7 +13,7 @@ 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``. +shardAccept.cu) and reached through the ``ops/spec`` entry ``m3_shard_accept``. """ from dataclasses import dataclass @@ -21,10 +21,10 @@ import torch -from .. import dispatch +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 +SHARD_ROWS = 48 # capacity: fixed: DECODE_MAX_TOKENS verify rows @dataclass(eq=False) @@ -107,7 +107,7 @@ def _workspace(mapping, device) -> Optional[AcceptWorkspace]: if _WS is not None: return _WS if torch.cuda.is_current_stream_capturing(): - dispatch.capture_missed(dispatch.SHARD_ACCEPT.path, None, "the shard-accept workspace") + # 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 @@ -154,7 +154,7 @@ def shard_accept( 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( + _shard_accept_op( logits, vocab_offset, ws.keys, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_spec_tail.py index 996371370fde..0d4dc9cd1b64 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_spec_tail.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_spec_tail.py @@ -13,8 +13,13 @@ 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}. +# 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( @@ -33,10 +38,9 @@ def draft_input( 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( + return _draft_input_op( ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, use_pdl ) - return out, residual def chunk_norm( @@ -45,8 +49,7 @@ def chunk_norm( """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 + 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 / @@ -72,7 +75,7 @@ def step_gather( and kv_off.is_contiguous() ): return False - torch.ops.trtllm.minimax_m3_step_gather( + _step_gather_op( new_tok, next_draft, lens, slots, ppi, P, T, D, ids, ids0, draft, d0, pos_off, kv_off, b0 ) return True @@ -88,7 +91,7 @@ def slot_scatter( 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( + _slot_scatter_op( slots, src_new, dst_new, @@ -110,4 +113,4 @@ def capture_norm( 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) + _capture_norm_op(h, residual, weight, eps, dst) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mx_gemv/__init__.py similarity index 100% rename from tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mx_gemv/__init__.py diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mx_gemv/m3_mx_gemv_kernel.py similarity index 99% rename from tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/m3_mx_gemv_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mx_gemv/m3_mx_gemv_kernel.py index b1c9d4de0e08..e3d7751231ee 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/m3_mx_gemv_kernel.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mx_gemv/m3_mx_gemv_kernel.py @@ -40,11 +40,13 @@ from cutlass import dsl_user_op from cutlass.experimental import primitives as prims -CTA_M = 128 # weight rows per CTA = the MMA's M +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 +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 +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 diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mx_gemv/op.py similarity index 99% rename from tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mx_gemv/op.py index abcbadff4721..d5440088cbc8 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mx_gemv/op.py @@ -28,14 +28,14 @@ from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl -MAX_TOKENS = 48 +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 +_SMEM_LIMIT = 220 * 1024 # capacity: fixed: the device's shared-memory limit def _mma_n(m: int) -> int: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mxfp8_low_m_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mxfp8_low_m_gemm.py index 04e21b99471e..cc635f93a836 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mxfp8_low_m_gemm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mxfp8_low_m_gemm.py @@ -38,7 +38,7 @@ (6144, 3072): (4, 2, 8, 2), # dense MLP down } # Opt-in shared memory per CTA on SM100 (227 KiB) less a margin for the static mbarriers. -_MAX_DYNAMIC_SMEM = 227 * 1024 - 1024 +_MAX_DYNAMIC_SMEM = 227 * 1024 - 1024 # capacity: fixed: the device's shared-memory limit _logged_shapes = set() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/triton_sparse_decode.py index c1e0f4332b88..064a3c0b2b10 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/triton_sparse_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/triton_sparse_decode.py @@ -57,7 +57,7 @@ from . import m3_attn_lc, m3_prequant # One sparse block is exactly one KV page. -SPARSE_BLOCK_SIZE = 128 +SPARSE_BLOCK_SIZE = 128 # capacity: fixed: the MSA block size of the checkpoint _FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) @@ -68,7 +68,7 @@ } # Total CTAs the split-K partitioning aims for before it stops splitting. -_TARGET_GRID = 256 +_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. @@ -77,7 +77,7 @@ # 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 +MERGE_MXFP8_MIN_TOKENS = 1 # capacity: fixed: a row lower bound class MergeMxfp8: @@ -89,7 +89,7 @@ class MergeMxfp8: merge rewrites it, in stream order. """ - MAX_ROWS = 128 + MAX_ROWS = 128 # capacity: fixed: the merge hand-off's rows (decode steps) def __init__(self): self._buffers = {} @@ -1176,7 +1176,7 @@ def _check_nvfp4_inputs( # 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 +MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query producer's crossover # Warps per multi-query producer program. MQ_NUM_WARPS = 8 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/layers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/layers.py deleted file mode 100644 index 3bfafc82cee6..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/layers.py +++ /dev/null @@ -1,1682 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""MiniMax-M3 NVFP4 / sm_103 / tp4: the decoder layers, on the built-in modules. - -The attention and MoE modules subclass the built-in ``MiniMaxM3Attention`` and -``MiniMaxM3MoE``: the generic path (prefill, mixed steps, wide steps) runs their -code unchanged, and the checkpoint loads through the built-in loader. What this -target adds is the decode-step kernels, each taken only by the step shapes it -was written for and measured on (GB300, tp4, 1-8 requests): - -* attention: for a short-sequence generation step of an NVFP4 sparse layer (every - request within 16 blocks, so the block selection is the identity) the packed - QKV GEMM is followed by one fused norm/RoPE/KV-insert/sparse-decode/merge - kernel; one request decoding one token runs QKV GEMV + attention + o_proj as - one launch, and one request verifying a draft runs attention + o_proj as one; -* MoE: the router (an fp32 GEMV or a tensor-core router fused with the routed - experts' NVFP4 activation quantization) and the routed experts as one engine - launch for 1-8 tokens, with the shared expert folded in; a second engine - covers 9-13 rows; -* the layer boundaries: a multicast all-reduce fused with the residual add and - the next RMSNorm up to 48 tokens. - -The decoder layer and the text model are this target's own: they construct the -subclassed modules. Their generic forward is the built-in one. -""" - -from __future__ import annotations - -from typing import Callable, Dict, Optional, Tuple - -import torch -from torch import nn -from transformers import PretrainedConfig - -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.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention -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_minimaxm3 import ( - MiniMaxM3Attention, - MiniMaxM3MoE, - _build_swiglu_oai_dense_mlp, - _nvfp4_cache_supports_fused_write, - _validate_sparse_attention_runtime_config, - get_moe_layer_ids, - get_sparse_disable_index_value_layer_ids, - get_sparse_layer_ids, - maybe_bcg_minimax_m3_attn_custom_op_inplace, -) -from tensorrt_llm._torch.models.modeling_utils import DecoderModel -from tensorrt_llm._torch.modules.decoder_layer import DecoderLayer -from tensorrt_llm._torch.modules.embedding import Embedding -from tensorrt_llm._torch.modules.linear import ( - flashinfer_mxfp8_autotune_active, - flashinfer_mxfp8_decode_graph_capture_active, -) -from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel -from tensorrt_llm._torch.modules.rms_norm import RMSNorm -from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph -from tensorrt_llm._torch.speculative import SpecMetadata -from tensorrt_llm._torch.utils import ( - AuxStreamType, - EventType, - Fp4QuantizedTensor, - is_torch_compiling, -) -from tensorrt_llm.logger import logger - -from . import dispatch -from .dispatch import ( - FUSED_ATTN_MAX_QUERY_LEN, - FUSED_ATTN_MAX_TOKENS, - MOE16_MAX_TOKENS, - MOE16_MIN_TOKENS, - MOE_ENGINE_MAX_TOKENS, - PRODUCER_EARLY_TRIGGER, - ROUTER_TC_MAX_TOKENS, -) -from .kernels import m3_mc_allreduce, m3_prequant -from .kernels.m3_router_gemv import m3_router_gemv, m3_router_gemv_fp4, m3_router_tc -from .linears import PUSH_ATTN, PUSH_MOE, mark_m3_linear -from .ops.attention.m3_index_score import SPLIT_K, m3_index_score -from .ops.attention.m3_producer import m3_producer -from .ops.moe.m3_moe_engine import m3_moe_engine -from .spec_worker import capture_hidden_states - -# The row ranges of the engines, the router and the fused attention are in dispatch.py. -# The shared expert folds into the routed engine from one token (m3_moe with dedicated shared-expert CTAs). -MOE_FOLD_MIN_TOKENS = 1 -# SMs left to the concurrent shared-expert kernels while a routed engine without the shared fold runs (persistent). -MOE_ENGINE_RESERVED_SMS = 24 -# CTAs per request of the CuTe DSL index scorer at every batch (each block is scored once, so the scores do not depend -# on it), and the request counts up to which it compiles the early-issue schedule (kernels/m3_index_score). -INDEX_SCORE_SPLIT_K = SPLIT_K -INDEX_SCORE_EARLY_MAX_REQUESTS = dispatch.INDEX_SCORE_EARLY_MAX_REQUESTS -# The generic layer's boundaries (prefill, Eagle3 capture layers): the low-latency multicast all-reduce + residual + -# RMSNorm beats the native one-shot all-reduce up to 8 tokens; with the shared + routed add folded in, up to its own -# maximum. -MC_AR_MAX_TOKENS = 8 - -# 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] = {} -# The CuTe DSL index scorer of the sparse decode (its kernels are compiled by prebuild_index_scorer). -INDEX_SCORE_PATH = dispatch.INDEX_SCORE.path - - -def _scorer_runner(): - from .kernels.m3_index_score import IndexScoreRunner - - return IndexScoreRunner - - -def _scorer_key( - dtype: torch.dtype, num_heads: int, query_len: int, head_dim: int, split_k: int, early: bool -) -> tuple: - """The scorer's compile key (``IndexScoreRunner``).""" - return (dtype, int(num_heads), int(query_len), int(head_dim), int(split_k), bool(early)) - - -def scorer_variants(layers, query_lens) -> list: - """The scorer keys a decode step of these layers can take: every sparse layer's (index dtype, index heads per - rank, index dim) at each decode query length, with the early-issue schedule (up to - ``INDEX_SCORE_EARLY_MAX_REQUESTS`` requests) and without it, where the kernel serves the geometry.""" - runner = _scorer_runner() - keys = [] - for layer in layers: - attn = getattr(layer, "self_attn", None) - if ( - attn is None - or not getattr(attn, "is_sparse_attention_layer", False) - or not isinstance(attn.attn, MiniMaxM3MsaSparseAttention) - ): - continue - config = attn.attn.m3_config - dtype = torch.float8_e4m3fn if attn.attn.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 = _scorer_key( - dtype, - config.num_index_heads, - qlen, - config.sparse_index_dim, - INDEX_SCORE_SPLIT_K, - early, - ) - if key not in keys: - keys.append(key) - return keys - - -def prebuild_index_scorer(layers, query_lens, device) -> list: - """Compile every key of ``scorer_variants`` 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 = _scorer_runner() - keys = scorer_variants(layers, 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 - - -_EXPERT_TENSOR_NAMES = ( - "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", -) - - -def _nvfp4_swiglu_oai_experts(experts: nn.Module): - """The routed experts' operands the engines read, or raise if the backend is not TRTLLM-Gen NVFP4 SwiGLU-OAI.""" - backend = getattr(experts, "backend", experts) - tensors = [getattr(backend, name, None) for name in _EXPERT_TENSOR_NAMES] - scale = getattr(backend, "fc31_input_scale", None) - if ( - any(t is None for t in tensors) - or scale is None - or getattr(backend, "scaling_vector_size", 16) != 16 - or getattr(backend, "fc31_act_scale", None) is not None - or not type(backend).quantize_input.__qualname__.startswith("TRTLLMGenNvfp4Base") - ): - raise ValueError("experts backend is not TRTLLM-Gen NVFP4 SwiGLU-OAI") - tensors = [t.data if isinstance(t, nn.Parameter) else t for t in tensors] - scale = scale.data if isinstance(scale, nn.Parameter) else scale - return tensors, scale - - -def _dense_mlp(model_config: ModelConfig, intermediate_size: int, layer_idx: int) -> nn.Module: - """The built-in dense SwiGLU-OAI MLP, its projections on this target's decode GEMM dispatch.""" - mlp = _build_swiglu_oai_dense_mlp( - model_config=model_config, intermediate_size=intermediate_size, layer_idx=layer_idx - ) - mark_m3_linear(mlp.gate_up_proj, low_m_gemm=True) - mark_m3_linear(mlp.down_proj, low_m_gemm=True) - return mlp - - -class M3MoE(MiniMaxM3MoE): - """MiniMax-M3 MoE with the decode-step router and expert engines. - - Engines are built per (token count, fusion) by ``prebuild_engines`` from - ``build_targets``, before the executor's warmup (building compiles and warms - the kernel); a forward only looks them up, and a step that finds no engine - is a miss and runs the built-in path. - """ - - def __init__( - self, - model_config: ModelConfig, - aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - layer_idx: int, - ) -> None: - super().__init__(model_config, aux_stream_dict, layer_idx) - self.layer_idx = layer_idx - if self.shared_experts is not None: - mark_m3_linear(self.shared_experts.gate_up_proj, low_m_gemm=False) - mark_m3_linear(self.shared_experts.down_proj, low_m_gemm=False, push=PUSH_MOE) - # The routed experts' NVFP4 input scale once resolved; False when they cannot take the fused router. - self._fp4_input_scale: Optional[torch.Tensor | bool] = None - self._engines: Dict[Tuple[int, bool, bool], Optional[object]] = {} - self._moe16_engines: Dict[int, Optional[object]] = {} - - def _router_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: - """The fp32 router logits: the fp32 GEMV at one row, the tensor-core router up to 48.""" - rows = hidden_states.shape[0] if hidden_states.dim() == 2 else 0 - path = dispatch.router_path(rows) - if path == dispatch.ROUTER_GEMV.path: - logits = m3_router_gemv(hidden_states, self.gate.weight, use_pdl=get_env_enable_pdl()) - if logits is not None: - return logits - dispatch.missed(path, rows, "the router GEMV declined the operands") - elif path == dispatch.ROUTER_TC.path: - fused = m3_router_tc(hidden_states, self.gate.weight, use_pdl=get_env_enable_pdl()) - if fused is not None: - return fused[0] - dispatch.missed(path, rows, "the tensor-core router declined the operands") - else: - dispatch.uncovered(dispatch.ROUTER_TC.path, f"more than {ROUTER_TC_MAX_TOKENS} rows") - return self.gate(hidden_states) - - def _router_fp4(self, hidden_states: torch.Tensor): - """Router logits and the routed experts' NVFP4 activation from one kernel, or None (then - ``_router_logits``).""" - if hidden_states.dim() != 2 or hidden_states.shape[0] > ROUTER_TC_MAX_TOKENS: - return None - if self._fp4_input_scale is None: - try: - _, scale = _nvfp4_swiglu_oai_experts(self.experts) - weight = getattr(self.experts, "backend", self.experts).w3_w1_weight - self._fp4_input_scale = ( - scale if weight.shape[-1] * 2 == hidden_states.shape[-1] else False - ) - except ValueError: - self._fp4_input_scale = False - logger.info( - f"[M3 router FP4] fused router+NVFP4 activation quant: {self._fp4_input_scale is not False}" - ) - rows = hidden_states.shape[0] - path = dispatch.router_path(rows) - if path == dispatch.ROUTER_GEMV.path: - path = dispatch.ROUTER_GEMV_FP4.path - if self._fp4_input_scale is False: - dispatch.missed( - path, - rows, - "the routed experts are not TRTLLM-Gen NVFP4 SwiGLU-OAI on this hidden size", - ) - return None - if path == dispatch.ROUTER_TC.path: - fused = m3_router_tc( - hidden_states, self.gate.weight, self._fp4_input_scale, use_pdl=get_env_enable_pdl() - ) - else: - fused = m3_router_gemv_fp4( - hidden_states, self.gate.weight, self._fp4_input_scale, use_pdl=get_env_enable_pdl() - ) - if fused is None: - dispatch.missed(path, rows, "the router declined the operands") - return fused - - def engine(self, num_tokens: int, shared: bool) -> Optional[object]: - """This layer's prebuilt routed-expert engine for ``num_tokens`` (1..8), router folded in. - - ``shared`` folds the shared expert in as well (routed + shared out). None - when the built-in path must run, after a miss: the engine failed to build, - or ``prebuild_engines`` did not build it (a forward never builds one). - """ - key = (num_tokens, True, shared) - return self._lookup(self._engines, key, dispatch.MOE_ENGINE.path, num_tokens) - - def moe16_engine(self, num_tokens: int) -> Optional[object]: - """This layer'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 or self.shared_experts is None: - return None - return self._lookup(self._moe16_engines, num_tokens, dispatch.MOE16_ENGINE.path, num_tokens) - - def _lookup(self, engines: dict, key, path: str, num_tokens: int) -> Optional[object]: - if key not in engines: - what = f"layer {self.layer_idx}'s engine" - if torch.cuda.is_current_stream_capturing(): - dispatch.capture_missed(path, num_tokens, what) - else: - dispatch.missed(path, num_tokens, f"{what} was not built before warmup") - return None - engine = engines[key] - if engine is None: - dispatch.missed(path, num_tokens, "its build failed (logged)") - return engine - - def prebuild_engines(self) -> None: - """Build (compile, on first use of each shared state, and warm) every engine a decode step of this layer can - take: ``MOE_ENGINE``'s token counts with the shared expert folded in, and ``MOE16_ENGINE``'s. Runs from - ``build_targets`` once the weights are real, on the current stream, in the same order on every rank; a build - that fails is logged and its token count falls back (a miss) in the forward. Whether an engine is kept is - agreed across the TP group (one host collective per engine, here): a build that fails on one rank only drops - that engine on every rank, so the ranks never mix an engine (and its boundary push) with the built-in MoE.""" - if self.shared_experts is None: - return - for num_tokens in sorted(dispatch.MOE_ENGINE.keys): - key = (num_tokens, True, True) - if key not in self._engines: - self._engines[key] = self._agreed( - dispatch.MOE_ENGINE.path, - num_tokens, - self._try_build( - dispatch.MOE_ENGINE.path, num_tokens, self._build_engine, shared=True - ), - ) - for num_tokens in sorted(dispatch.MOE16_ENGINE.keys): - if num_tokens not in self._moe16_engines: - self._moe16_engines[num_tokens] = self._agreed( - dispatch.MOE16_ENGINE.path, - num_tokens, - self._try_build( - dispatch.MOE16_ENGINE.path, num_tokens, self._build_moe16_engine - ), - ) - - def _agreed(self, path: str, num_tokens: int, engine: Optional[object]) -> Optional[object]: - return m3_mc_allreduce.tp_agree_built( - self.mapping, - engine, - f"[M3 dispatch] {path} at {num_tokens} rows: layer {self.layer_idx}'s engine", - ) - - def _try_build(self, path: str, num_tokens: int, build, **kwargs) -> Optional[object]: - try: - return build(num_tokens, **kwargs) - except (ImportError, RuntimeError, ValueError) as err: - logger.warning( - f"[M3 dispatch] {path} at {num_tokens} rows: layer {self.layer_idx}'s engine failed to build: {err}" - ) - return None - - def _build_engine(self, num_tokens: int, shared: bool) -> object: - from .ops.moe import m3_moe_engine as m3_moe_op - - tensors, scale = _nvfp4_swiglu_oai_experts(self.experts) - device = tensors[0].device - sms = torch.cuda.get_device_properties(device).multi_processor_count - kwargs = {} - if shared: - gu, dn = self.shared_experts.gate_up_proj, self.shared_experts.down_proj - kwargs["shared_weights"] = ( - gu.weight.data.view(torch.float8_e4m3fn), - gu.weight_scale.data, - dn.weight.data.view(torch.float8_e4m3fn), - dn.weight_scale.data, - ) - state = _ENGINE_STATES.get((str(device), num_tokens, shared)) - if state is None: - state = _ENGINE_STATES[(str(device), num_tokens, shared)] = m3_moe_op.M3MoeState( - device, - num_tokens=num_tokens, - num_ctas=sms - MOE_ENGINE_RESERVED_SMS, - shared=shared, - ) - engine = state.layer( - *tensors, - self.gate.e_score_correction_bias.data, - router_weight=self.gate.weight.data, - fc31_input_scale=scale.reshape(1).float(), - **kwargs, - ) - # Compile (first use of the state) and warm. - engine( - torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device) - ) - return engine - - def _build_moe16_engine(self, num_tokens: int) -> object: - from .ops.moe import m3_moe16_engine as op16 - - tensors, scale = _nvfp4_swiglu_oai_experts(self.experts) - device = tensors[0].device - gu, dn = self.shared_experts.gate_up_proj, self.shared_experts.down_proj - state = _MOE16_STATES.get((str(device), num_tokens)) - if state is None: - state = _MOE16_STATES[(str(device), num_tokens)] = op16.M3Moe16State( - device, num_tokens=num_tokens - ) - engine = state.layer( - *tensors, - self.gate.e_score_correction_bias.data, - router_weight=self.gate.weight.data, - fc31_input_scale=scale.reshape(1).float(), - shared_weights=( - gu.weight.data.view(torch.float8_e4m3fn), - gu.weight_scale.data, - dn.weight.data.view(torch.float8_e4m3fn), - dn.weight_scale.data, - ), - ) - # Compile (first use of the state) and warm. - engine( - torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device) - ) - return engine - - def forward( - self, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - final_all_reduce_params: Optional[AllReduceParams] = None, - lora_params: Optional[dict] = None, - do_finalize: bool = True, - return_parts: bool = False, - ): - """The MoE output; or, with ``return_parts``, ``(shared, routed)`` (routed None when the - engine folded the shared expert in); or, with ``do_finalize=False``, the unreduced - ``[shared_output, fc2_output, expert_scale_factor, expanded_idx_to_permuted_idx]`` - for the fused finalize + all-reduce + residual + RMSNorm epilogue. - - Mirrors the built-in ``MiniMaxM3MoE.forward`` apart from the decode engines. - """ - all_rank_num_tokens = attn_metadata.all_rank_num_tokens - engine_rows = ( - do_finalize - and hidden_states.dim() == 2 - and hidden_states.dtype == torch.bfloat16 - and 1 <= hidden_states.shape[0] <= MOE_ENGINE_MAX_TOKENS - ) - - def _compute_routed_output(): - fused = self._router_fp4(hidden_states) - if fused is not None: - router_logits, x_fp4, x_sf = fused - return self.experts( - Fp4QuantizedTensor(x_fp4, x_sf, is_sf_swizzled=False), - router_logits, - do_finalize=do_finalize, - output_dtype=hidden_states.dtype, - all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=False, - lora_params=lora_params, - ) - router_logits = self._router_logits(hidden_states) - return self.experts( - hidden_states, - router_logits, - do_finalize=do_finalize, - all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=False, - lora_params=lora_params, - ) - - def _compute_shared_output(): - return self.shared_experts(hidden_states, lora_params=lora_params) - - folded = None - if ( - engine_rows - and self.shared_experts is not None - and hidden_states.shape[0] >= MOE_FOLD_MIN_TOKENS - ): - engine = self.engine(hidden_states.shape[0], shared=True) - if engine is not None: - arp = m3_prequant.MOE_PUSH.take_engine() - folded = m3_moe_engine(hidden_states.contiguous(), engine, arp=arp) - if arp is not None: - m3_prequant.MOE_PUSH.engine_pushed = folded.data_ptr() - if folded is not None: - # Routed + shared from one kernel. - if return_parts: - return folded, None - result = folded - elif self.shared_experts is None: - result = _compute_routed_output() - else: - routed_output, shared_output = maybe_execute_in_parallel( - _compute_routed_output, - _compute_shared_output, - self.event_dict[EventType.Main], - self.event_dict[EventType.MoeShared], - # 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 self.aux_stream, - disable_on_compile=True, - ) - if not do_finalize: - return [shared_output, *routed_output] - if return_parts: - return shared_output, routed_output - # In-place add into ``shared_output`` to avoid allocating a temporary. - result = shared_output.add_(routed_output) - # tp4 routed experts are tensor-parallel (never fused-communication), so the sum is still rank-local. - return self.allreduce(result, all_reduce_params=final_all_reduce_params) - - -def producer_pw_covered(layer_uses_nvfp4, num_heads_q, num_heads_kv, num_heads_index) -> bool: - """Whether ``dispatch.PRODUCER`` covers a sparse layer: an NVFP4 KV cache at the TP4 heads (16 query, one KV and - one index head). A covered layer that cannot take ``m3_producer`` is a miss.""" - return bool(layer_uses_nvfp4) and (num_heads_q, num_heads_kv, num_heads_index) == (16, 1, 1) - - -def _producer_miss(covered: bool, rows: int, reason: str) -> None: - """A producer precondition failed: a miss of ``dispatch.PRODUCER`` on a covered layer (raised under - ``modeling_v2='require'``), otherwise nothing; the caller then runs the unfused producer.""" - if covered: - dispatch.missed(dispatch.PRODUCER.path, rows, reason) - return None - - -def producer_pw_supported(num_heads_q, num_heads_kv, num_heads_index, kv_data, kv_scale) -> bool: - """Whether a sparse layer's NVFP4 producer can take ``m3_producer``: the TP4 heads (16 query, one KV and one - index head) and the NVFP4 pool [pages, 2, 1, 128, 64] with its scale pool [pages, 2, 1, 128, 8].""" - return ( - num_heads_q == 16 - and num_heads_kv == 1 - and num_heads_index == 1 - and 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] - ) - - -class M3Attention(MiniMaxM3Attention): - """MiniMax-M3 attention with the decode-step fused kernels.""" - - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - mark_m3_linear(self.qkv_proj, low_m_gemm=True) - mark_m3_linear(self.o_proj, low_m_gemm=True, push=PUSH_ATTN) - self._mega_workspace = None - self._c1_workspace = None - if not self.is_sparse_attention_layer: - # The KV split trtllm-gen dense decode runs at, per request count (msa_decode.dense_decode_max_seq_len). - self.attn.dense_decode_max_seq_len_by_batch = dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH - - def _check_producer_projection(self) -> None: - """A layer the producer covers (dispatch.PRODUCER) needs the fused QKV/index projection: without it the step is - a miss (raises under require) instead of silently taking the unfused path.""" - if not self.enable_fused_qkv_index_projection and producer_pw_covered( - self.main_kv_is_nvfp4, - self.num_heads, - self.num_key_value_heads, - self.sparse_num_index_heads, - ): - dispatch.missed(dispatch.PRODUCER.path, None, "fuse_qkv_index_projection is disabled") - - def _fused_fp8_qkv_indexer_norm_rope_kv_insert( - self, - packed: torch.Tensor, - position_ids: Optional[torch.Tensor], - attn_metadata: AttentionMetadata, - *, - cache_tensors: Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = None, - ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: - """Run the vLLM-style horizontal producer for every sparse batch. - - Mirrors the built-in ``MiniMaxM3Attention._fused_fp8_qkv_indexer_norm_rope_kv_insert``; an NVFP4 cache takes - this target's ``m3_producer`` (the native producer bit for bit, its constant loads staged before the grid - wait and its dependents triggered on entry) where its geometry is the one that kernel is built for. - - The CUDA kernel is token-major and batch-type agnostic: per-token - positions and cache slots cover pure prefill, mixed aggregate batches, - and CUDA-graph decode. FP8 caches receive FP8 K/V directly. NVFP4 - caches receive packed E2M1 K/V plus their E4M3 SF16 scale bytes while - Q, index-Q, and index-K remain FP8. - """ - if ( - not self.enable_fused_qkv_index_projection - or not isinstance(self.attn, MiniMaxM3MsaSparseAttention) - or not (self.main_kv_is_fp8 or self.main_kv_is_nvfp4) - or self.attn.indexer_kv_dtype != "fp8" - ): - return None - # A sparse layer on an NVFP4 cache at the TP4 heads is the producer's (dispatch.PRODUCER): every precondition - # below that fails for it is a miss. - covered = producer_pw_covered( - self.main_kv_is_nvfp4, - self.num_heads, - self.num_key_value_heads, - self.sparse_num_index_heads, - ) - rows = int(packed.shape[0]) - rope = self.pos_embd_params.rope if self.pos_embd_params is not None else None - rotary_dim = int(rope.dim) if rope is not None else 0 - if ( - packed.dtype != torch.bfloat16 - or position_ids is None - or self.head_dim != 128 - or self.sparse_index_dim != 128 - or self.sparse_num_index_heads != self.num_key_value_heads - or not self.use_gemma_norm - or self.pos_embd_params is None - or not self.pos_embd_params.is_neox - or self.rotary_emb is None - or rope is None - or rotary_dim != 64 - ): - return _producer_miss( - covered, - rows, - "the packed QKV dtype, position ids, head dims, norm or RoPE configuration", - ) - norm_eps = self.q_norm.variance_epsilon - if any( - module.variance_epsilon != norm_eps - for module in (self.k_norm, self.index_q_norm, self.index_k_norm) - ): - return _producer_miss(covered, rows, "the four RMSNorm epsilons differ") - norm_weights = ( - self.q_norm.weight, - self.k_norm.weight, - self.index_q_norm.weight, - self.index_k_norm.weight, - ) - if any(weight.dtype != torch.bfloat16 or not weight.is_cuda for weight in norm_weights): - return _producer_miss(covered, rows, "a norm weight is not BF16 on the GPU") - - kv_cache_manager = attn_metadata.kv_cache_manager - if cache_tensors is None: - if kv_cache_manager is None: - return _producer_miss(covered, rows, "the step has no KV cache manager") - buffers = kv_cache_manager.get_buffers(self.layer_idx, kv_layout="HND") - index_k_cache = attn_metadata.msa_idx_k_cache(self.layer_idx) - out_cache_loc = getattr(attn_metadata, "msa_out_cache_loc", None) - else: - buffers, index_k_cache, out_cache_loc = cache_tensors - num_tokens = int(packed.shape[0]) - layer_uses_nvfp4 = bool( - getattr(kv_cache_manager, "is_nvfp4_layer", lambda _layer_idx: False)(self.layer_idx) - ) - scale_buffers = ( - kv_cache_manager.get_block_scale_buffers(self.layer_idx, kv_layout="HND") - if layer_uses_nvfp4 - else None - ) - kv_quant_scale = getattr(self.qkv_proj, "inv_kv_scales", None) if layer_uses_nvfp4 else None - if layer_uses_nvfp4: - supports_fused_layout = ( - buffers is not None - and scale_buffers is not None - and _nvfp4_cache_supports_fused_write( - buffers, scale_buffers, self.num_key_value_heads - ) - ) - supported_main_cache = ( - supports_fused_layout - and buffers.is_cuda - and scale_buffers.is_cuda - and scale_buffers.device == buffers.device - and kv_quant_scale is not None - 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_cache = ( - buffers is not None - and buffers.is_cuda - and buffers.dtype == torch.float8_e4m3fn - and buffers.dim() == 5 - and tuple(buffers.shape[1:]) == (2, self.num_key_value_heads, 128, 128) - and buffers.stride(4) == 1 - and buffers.stride(3) == 128 - and buffers.stride(2) == 128 * 128 - and buffers.stride(1) >= self.num_key_value_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_cache = ( - 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] - ) - rotary_cos_sin = self.rotary_emb.rotary_cos_sin - supported_rope_cache = ( - rotary_cos_sin.is_cuda - and rotary_cos_sin.dtype == torch.float32 - and rotary_cos_sin.is_contiguous() - and rotary_cos_sin.dim() == 3 - and tuple(rotary_cos_sin.shape[1:]) == (2, 32) - ) - if ( - not supported_main_cache - or not supported_index_cache - or not supported_rope_cache - 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() < num_tokens - ): - return _producer_miss( - covered and layer_uses_nvfp4, - num_tokens, - "the layer's NVFP4 / index caches, RoPE table, KV quant scale or cache slots are not the " - "layout the producer is built for", - ) - - common_args = ( - out_cache_loc[:num_tokens], - self.num_heads, - self.num_key_value_heads, - self.sparse_num_index_heads, - self.head_dim, - rotary_dim, - norm_eps, - self.q_norm.weight, - self.k_norm.weight, - self.index_q_norm.weight, - self.index_k_norm.weight, - rotary_cos_sin, - position_ids.reshape(-1).contiguous().to(torch.int32), - ) - if layer_uses_nvfp4 and producer_pw_supported( - self.num_heads, - self.num_key_value_heads, - self.sparse_num_index_heads, - buffers, - scale_buffers, - ): - q, index_q = m3_producer( - packed.contiguous(), - buffers.view(torch.uint8), - scale_buffers.view(torch.uint8), - index_k_cache, - common_args[0], - kv_quant_scale, - norm_eps, - *common_args[7:12], - common_args[12], - self.num_heads, - get_env_enable_pdl(), - PRODUCER_EARLY_TRIGGER, - ) - elif layer_uses_nvfp4: - if covered: - dispatch.missed( - dispatch.PRODUCER.path, - num_tokens, - f"the NVFP4 pool {tuple(buffers.shape)} / scale pool " - f"{tuple(scale_buffers.shape)} is not the producer's layout", - ) - q, index_q = torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( - packed.contiguous(), - buffers.view(torch.uint8), - scale_buffers, - index_k_cache, - common_args[0], - kv_quant_scale, - *common_args[1:], - ) - else: - q, index_q = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( - packed.contiguous(), - buffers, - index_k_cache, - *common_args, - ) - return q.flatten(1), index_q.flatten(1) - - def _fused_attn_args(self, position_ids, attn_metadata) -> Optional[dict]: - """Arguments of the fused decode attention for this step, or None when it does not apply. - - The layer must be an NVFP4 sparse MSA layer and the step a pure generation - step whose selection the metadata staged as the identity (every request - within 16 blocks of 128 tokens), with at most 8 query tokens per request. - """ - path = dispatch.FUSED_DECODE_ATTN.path - if ( - not isinstance(self.attn, MiniMaxM3MsaSparseAttention) - or not self.is_sparse_attention_layer - or not self.main_kv_is_nvfp4 - or self.attn.indexer_kv_dtype != "fp8" - ): - dispatch.uncovered(path, "not an NVFP4 sparse MSA layer with an FP8 indexer") - return None - if position_ids is None or is_torch_compiling() or attn_metadata.num_contexts > 0: - dispatch.uncovered(path, "not a pure generation step with positions") - return None - if not attn_metadata._msa_short_ok: - dispatch.uncovered( - path, "a request beyond 16 KV blocks: the block selection is not the identity" - ) - return None - span = attn_metadata.msa_decode_span - if span is None or span.row_first != 0: - dispatch.uncovered(path, "no decode span starting at row 0") - return None - num_reqs = int(attn_metadata.num_seqs) - qlen = int(span.query_len) - if not dispatch.FUSED_DECODE_ATTN.covers((qlen, num_reqs)): - dispatch.uncovered( - path, - f"more than {FUSED_ATTN_MAX_QUERY_LEN} query rows per request or " - f"{FUSED_ATTN_MAX_TOKENS} per step", - ) - return None - manager = attn_metadata.kv_cache_manager - if manager is not None and not manager.is_nvfp4_layer(self.layer_idx): - dispatch.uncovered(path, "the layer's KV cache is not NVFP4") - return None - from .kernels import m3_fused_attn - - if m3_fused_attn._module is None: - # Building compiles every geometry (minutes): only outside capture (the eager warmup steps). - if torch.cuda.is_current_stream_capturing(): - dispatch.capture_missed(path, num_reqs * qlen, "the fused attention module") - return None - m3_fused_attn._load_module() - kv_scales = getattr(self.qkv_proj, "kv_scales", None) - inv_kv_scales = getattr(self.qkv_proj, "inv_kv_scales", None) - if manager is None or kv_scales is None or inv_kv_scales is None: - dispatch.missed( - path, - num_reqs * qlen, - "no KV-cache manager" if manager is None else "the QKV projection has no KV scales", - ) - return None - k_scale, v_scale = _aligned_nvfp4_dequant_scales(self.attn, kv_scales) - return dict( - num_tokens=num_reqs * qlen, - decode_query_len=qlen, - kv_data=manager.get_buffers(self.layer_idx, kv_layout="HND").view(torch.uint8), - kv_scale=manager.get_block_scale_buffers(self.layer_idx, kv_layout="HND"), - index_k_cache=attn_metadata.msa_idx_k_cache(self.layer_idx), - out_cache_loc=attn_metadata.msa_out_cache_loc, - position_ids=position_ids.reshape(-1).contiguous().to(torch.int32), - block_table=attn_metadata.msa_block_table[:num_reqs], - seq_lens=attn_metadata.msa_seq_lens_cuda[:num_reqs], - kv_quant_scale=inv_kv_scales, - k_global_scale=k_scale, - v_global_scale=v_scale, - ) - - def _attn_mega(self, hidden_states: torch.Tensor, args: dict) -> Optional[torch.Tensor]: - """o_proj output of the one-launch QKV + attention + o_proj kernel for one request decoding one token whose - layer input the preceding all-reduce handed over in MXFP8; None, with nothing consumed, otherwise.""" - if not dispatch.ATTN_MEGA.covers( - (args["decode_query_len"], int(args["seq_lens"].shape[0])) - ): - return None - if not flashinfer_mxfp8_decode_graph_capture_active(): - return None - pending = m3_prequant.SLOT.pending - if ( - pending is None - or hidden_states.dim() != 2 - or pending[0] != hidden_states.data_ptr() - or pending[1] != 1 - or pending[2] != hidden_states.shape[1] - ): - # After a layer whose boundary does not quantize (a dense layer, an Eagle3 capture layer). - dispatch.uncovered(dispatch.ATTN_MEGA.path, "the layer input has no MXFP8 hand-off") - return None - if self.qkv_proj.weight.shape != (2560, 6144) or self.o_proj.weight.shape != (6144, 2048): - dispatch.uncovered(dispatch.ATTN_MEGA.path, "not the tp4 QKV / o_proj shapes") - return None - from .kernels import m3_attn_mega - - if self._mega_workspace is None: - # The first call JIT-builds the module; both only outside capture (the eager warmup steps). - if torch.cuda.is_current_stream_capturing(): - dispatch.capture_missed( - dispatch.ATTN_MEGA.path, 1, "the attn_mega module and workspace" - ) - return None - m3_attn_mega._load_module() - self._mega_workspace = m3_attn_mega.MegaWorkspace(hidden_states.device) - xq, xsf = m3_prequant.SLOT.take(hidden_states) - ar = m3_prequant.AR_PUSH.take_target() - output = hidden_states.new_empty( - (1, self.num_heads * self.head_dim), dtype=self.attn_activation_dtype - ) - y = hidden_states.new_empty((1, self.o_proj.weight.shape[0]), dtype=torch.bfloat16) - m3_attn_mega.m3_attn_mega( - xq, - xsf, - self.qkv_proj.weight, - self.qkv_proj.weight_scale, - self._mega_workspace, - args["kv_data"], - args["kv_scale"], - args["index_k_cache"], - args["out_cache_loc"], - args["position_ids"], - args["block_table"], - args["seq_lens"], - args["kv_quant_scale"], - args["k_global_scale"], - args["v_global_scale"], - self.q_norm.weight, - self.k_norm.weight, - self.index_k_norm.weight, - self.rotary_emb.rotary_cos_sin, - self.o_proj.weight, - self.o_proj.weight_scale, - eps=self.q_norm.variance_epsilon, - sm_scale=(self.head_dim**-0.5) / float(self.attn.q_scaling), - decode_query_len=1, - output=output, - y=y, - use_pdl=get_env_enable_pdl(), - ar=ar, - ) - if ar is not None: - m3_prequant.AR_PUSH.pushed = y.data_ptr() - return y - - def _attn_c1(self, packed: torch.Tensor, args: dict) -> Optional[torch.Tensor]: - """o_proj output of the one-launch attention + o_proj kernel for one request verifying a draft, pushed into - the attention boundary's all-reduce; None, with nothing consumed, when the step or the layer does not - qualify.""" - qlen, num_tokens = args["decode_query_len"], args["num_tokens"] - if not dispatch.ATTN_C1.covers((qlen, int(args["seq_lens"].shape[0]))): - return None - if not flashinfer_mxfp8_decode_graph_capture_active(): - return None - if m3_prequant.AR_PUSH.target is None: - dispatch.missed( - dispatch.ATTN_C1.path, num_tokens, "no armed attention-boundary push to write into" - ) - return None - if self.o_proj.weight.shape != (6144, 2048): - dispatch.uncovered(dispatch.ATTN_C1.path, "not the tp4 o_proj shape") - return None - from .kernels import m3_attn_mega_c1 - - if self._c1_workspace is None: - # The first call JIT-builds the module; both only outside capture (the eager warmup steps). - if torch.cuda.is_current_stream_capturing(): - dispatch.capture_missed( - dispatch.ATTN_C1.path, num_tokens, "the attn_c1 module and workspace" - ) - return None - m3_attn_mega_c1._load_module() - self._c1_workspace = m3_attn_mega_c1.MegaWorkspace(packed.device) - ar = m3_prequant.AR_PUSH.take_target() - output = packed.new_empty( - (num_tokens, self.num_heads * self.head_dim), dtype=self.attn_activation_dtype - ) - y = packed.new_empty((num_tokens, self.o_proj.weight.shape[0]), dtype=torch.bfloat16) - m3_attn_mega_c1.m3_attn_c1( - packed.contiguous(), - self._c1_workspace, - args["kv_data"], - args["kv_scale"], - args["index_k_cache"], - args["out_cache_loc"], - args["position_ids"], - args["block_table"], - args["seq_lens"], - args["kv_quant_scale"], - args["k_global_scale"], - args["v_global_scale"], - self.q_norm.weight, - self.k_norm.weight, - self.index_k_norm.weight, - self.rotary_emb.rotary_cos_sin, - self.o_proj.weight, - self.o_proj.weight_scale, - eps=self.q_norm.variance_epsilon, - sm_scale=(self.head_dim**-0.5) / float(self.attn.q_scaling), - decode_query_len=qlen, - output=output, - y=y, - use_pdl=get_env_enable_pdl(), - ar=ar[:5], - ) - m3_prequant.AR_PUSH.pushed = y.data_ptr() - return y - - def _fused_attn(self, packed: torch.Tensor, args: dict) -> torch.Tensor: - """The fused norm/RoPE/KV-insert/sparse-decode/merge kernel after the packed QKV GEMM.""" - from .kernels.m3_fused_attn import m3_fused_decode_attn - from .kernels.triton_sparse_decode import MERGE_MXFP8, MERGE_MXFP8_MIN_TOKENS - - num_tokens = args["num_tokens"] - width = self.num_heads * self.head_dim - output = packed.new_empty((num_tokens, width), dtype=self.attn_activation_dtype) - mxfp8 = None - if MERGE_MXFP8_MIN_TOKENS <= num_tokens <= MERGE_MXFP8.MAX_ROWS: - mxfp8 = MERGE_MXFP8.buffers(width, output.device) - m3_fused_decode_attn( - packed.contiguous(), - args["kv_data"], - args["kv_scale"], - args["index_k_cache"], - args["out_cache_loc"], - args["position_ids"], - args["block_table"], - args["seq_lens"], - args["kv_quant_scale"], - args["k_global_scale"], - args["v_global_scale"], - self.q_norm.weight, - self.k_norm.weight, - self.index_k_norm.weight, - self.rotary_emb.rotary_cos_sin, - eps=self.q_norm.variance_epsilon, - sm_scale=(self.head_dim**-0.5) / float(self.attn.q_scaling), - decode_query_len=args["decode_query_len"], - output=output, - mxfp8_out=mxfp8, - use_pdl=get_env_enable_pdl(), - ) - if mxfp8 is not None: - m3_prequant.SLOT.publish(output, mxfp8[0][:num_tokens], mxfp8[1]) - return output - - def _msa_attention_core(self, q, k, v, idx_q, idx_k, attn_metadata, output): - """The built-in MSA attention core, except a pure generation step of a sparse layer whose caches the - fused producer already wrote: that one runs this target's selection and sparse decode.""" - span = attn_metadata.msa_decode_span - if ( - self.is_sparse_attention_layer - and k is None - and attn_metadata.num_contexts == 0 - and span is not None - and span.row_first == 0 - ): - self._decode_sparse_core(q, idx_q, attn_metadata, output) - return output - return super()._msa_attention_core(q, k, v, idx_q, idx_k, attn_metadata, output) - - def _decode_sparse_core(self, q, idx_q, attn_metadata, output) -> None: - """Top-16 block selection and the block-sparse decode of one pure generation step. - - 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). - """ - from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( - MSA_REQUIRED_TOPK, - msa_paged_kv, - select_blocks_from_maxscore, - ) - - from .kernels.long_select import long_select_blocks - from .kernels.triton_sparse_decode import minimax_m3_sparse_attn_decode - - md = attn_metadata - attn = self.attn - config = attn.m3_config - num_tokens = int(q.shape[0]) - rows = int(md.num_seqs) - qlen = int(md.msa_decode_span.query_len) - block_table = md.msa_block_table[:rows] - seq_lens = md.msa_seq_lens_cuda[:rows] - num_kv_heads = int(attn.indexer.config.num_kv_heads) - # The index scorer's kernels are compiled and loaded by build_targets (prebuild_index_scorer), so no step - # compiles one and the identity can be taken from the first decode step. - identity = md._msa_short_ok and md.msa_short_table is not None - if identity: - topk = md.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 = md.msa_idx_k_cache(self.layer_idx) - max_score = md.msa_proxy_max_score_view( - config.num_index_heads, md.msa_worst_case_max_k_tiles, num_tokens - ) - early = rows <= INDEX_SCORE_EARLY_MAX_REQUESTS - key = _scorer_key( - idx_q_view.dtype, - config.num_index_heads, - qlen, - idx_q_view.shape[-1], - INDEX_SCORE_SPLIT_K, - early, - ) - if key not in _scorer_runner().kernel_cache: - # Outside require mode the op below compiles it (the only scorer this step has). - dispatch.missed( - INDEX_SCORE_PATH, - num_tokens, - f"the index scorer {key} was not compiled before warmup", - ) - 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, - ) - n_valid_blocks = md.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 - ) - 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=num_tokens <= dispatch.SELECT_EARLY_MAX_ROWS, - ) - if table is None: - dispatch.uncovered( - dispatch.LONG_SELECT.path, - f"{max_score_kv.shape[1]} KV blocks, outside {dispatch.LONG_SELECT_MIN_BLOCKS}.." - f"{dispatch.LONG_SELECT_MAX_BLOCKS}: the built-in selector", - ) - table = select_blocks_from_maxscore( - max_score_kv, - topk=MSA_REQUIRED_TOPK, - n_valid_blocks=n_valid_blocks, - init_blocks=config.init_blocks, - local_blocks=config.local_blocks, - head_major_output=True, - ) - topk = table.permute(1, 0, 2) - manager = md.kv_cache_manager - k_paged, v_paged = msa_paged_kv(manager, self.layer_idx) - nvfp4_args = {} - if manager.is_nvfp4_layer(self.layer_idx): - k_scale, v_scale = _aligned_nvfp4_dequant_scales(attn, self.qkv_proj.kv_scales) - scales = manager.get_block_scale_buffers(self.layer_idx, "HND") - nvfp4_args = dict( - k_block_scale=scales[:, 0], - v_block_scale=scales[:, 1], - k_global_scale=k_scale, - v_global_scale=v_scale, - ) - minimax_m3_sparse_attn_decode( - q.view(num_tokens, self.num_heads, self.head_dim), - k_paged, - v_paged, - topk, - block_table, - seq_lens, - sm_scale=(self.head_dim**-0.5) / float(attn.q_scaling), - output=output.view(num_tokens, self.num_heads, self.head_dim), - decode_query_len=qlen, - identity_blocks=identity, - **nvfp4_args, - ) - - def _dense_decode_fp8_q(self, attn_metadata) -> bool: - """On decode-only steps of a hybrid NVFP4 cache, the dense layers (whose KV pool is FP8) take FP8 q/k/v - straight from the fused norm/RoPE kernel instead of casting them separately.""" - return ( - self.main_kv_is_nvfp4 - and not self.is_sparse_attention_layer - and self._msa_backend_active() - and attn_metadata.num_contexts == 0 - ) - - def _dense_forward( - self, - position_ids: torch.IntTensor, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - all_reduce_params: Optional[AllReduceParams] = None, - **kwargs, - ) -> torch.Tensor: - """The built-in dense attention (layers 0-2), FP8 q/k/v on decode steps. - - Mirrors the built-in ``MiniMaxM3Attention._dense_forward``. - """ - if attn_metadata is None: - raise RuntimeError( - f"MiniMax-M3 dense forward (layer {self.layer_idx}) requires attn_metadata; received None." - ) - qkv = self.qkv_proj(hidden_states) - fused_qkv = self._fused_qk_norm_rope( - qkv, - position_ids, - num_heads_q=self.num_heads, - num_heads_k=self.num_key_value_heads, - num_heads_v=self.num_key_value_heads, - head_dim=self.head_dim, - q_norm=self.q_norm, - k_norm=self.k_norm, - out_fp8=self._emit_fp8_main_qkv() or self._dense_decode_fp8_q(attn_metadata), - ) - if fused_qkv is not None: - q, k, v = self._split_main_qkv(fused_qkv) - else: - assert not self._expect_fused_qk_norm_rope(position_ids), ( - f"MiniMax-M3 dense attention (layer {self.layer_idx}) expected the " - f"fused QK-norm+RoPE kernel (bf16 activations, head_dim=" - f"{self.head_dim}) but fell back to the separate path; qkv dtype " - f"is {qkv.dtype} (expected {self.attn_activation_dtype})." - ) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = self.apply_qk_norm(q, k) - if self.rotary_emb is not None and position_ids is not None: - q, k = self.rotary_emb(position_ids, [q, k]) - o = self._forward_attention_core(q, k, v, None, None, attn_metadata) - return self.o_proj(o, all_reduce_params=all_reduce_params) - - def _sparse_forward( - self, - position_ids: torch.IntTensor, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - all_reduce_params: Optional[AllReduceParams] = None, - **kwargs, - ) -> torch.Tensor: - """The built-in sparse attention, with the fused decode kernels on the steps they take. - - Mirrors the built-in ``MiniMaxM3Attention._sparse_forward``. - """ - if attn_metadata is None: - raise RuntimeError( - f"MiniMax-M3 sparse forward (layer {self.layer_idx}) requires attn_metadata; received None." - ) - packed_qkv = None - packed_idx_qk = None - self._check_producer_projection() - if self.enable_fused_qkv_index_projection: - if ( - self.register_to_config - and is_torch_compiling() - and isinstance(self.attn, MiniMaxM3MsaSparseAttention) - and self._emit_fp8_main_qkv() - and self.attn.indexer_kv_dtype == "fp8" - ): - kv_cache, index_k_cache = attn_metadata.msa_layer_cache_tensors[self.layer_idx] - q, idx_q = torch.ops.trtllm.minimax_m3_fused_sparse_qkv_producer( - hidden_states, - position_ids, - kv_cache, - index_k_cache, - attn_metadata.msa_out_cache_loc, - self.layer_idx_str, - ) - o = self._forward_attention_core(q, None, None, idx_q, None, attn_metadata) - return self.o_proj(o, all_reduce_params=all_reduce_params) - if self.register_to_config and (is_torch_compiling() or is_in_breakable_cuda_graph()): - packed = torch.ops.trtllm.minimax_m3_qkv_index_proj( - hidden_states, position_ids, self.layer_idx_str - ) - num_tokens = ( - position_ids.shape[-1] if position_ids is not None else hidden_states.shape[0] - ) - output = packed.new_empty( - (num_tokens, self.num_heads * self.head_dim), dtype=self.attn_activation_dtype - ) - maybe_bcg_minimax_m3_attn_custom_op_inplace( - None, None, None, None, None, packed, position_ids, self.layer_idx_str, output - ) - return self.o_proj(output, all_reduce_params=all_reduce_params) - # Gathered before the QKV GEMM so nothing runs between it and the fused attention (PDL hand-off). - fused_attn = self._fused_attn_args(position_ids, attn_metadata) - if fused_attn is not None: - y = self._attn_mega(hidden_states, fused_attn) - if y is not None: - return y - packed = self.qkv_proj(hidden_states) - if fused_attn is not None and ( - packed.shape[0] != fused_attn["num_tokens"] or packed.shape[1] != 2560 - ): - dispatch.missed( - dispatch.FUSED_DECODE_ATTN.path, - fused_attn["num_tokens"], - f"the packed QKV is {tuple(packed.shape)}, not [rows, 2560]", - ) - fused_attn = None - if fused_attn is not None: - y = self._attn_c1(packed, fused_attn) - if y is not None: - return y - o = self._fused_attn(packed, fused_attn) - return self.o_proj(o, all_reduce_params=all_reduce_params) - horizontal = self._fused_fp8_qkv_indexer_norm_rope_kv_insert( - packed, position_ids, attn_metadata - ) - if horizontal is not None: - q, idx_q = horizontal - o = self._forward_attention_core(q, None, None, idx_q, None, attn_metadata) - return self.o_proj(o, all_reduce_params=all_reduce_params) - main_size = self.q_size + 2 * self.kv_size - packed_qkv, packed_idx_qk = packed.split( - [main_size, self.index_q_size + self.index_k_size], dim=-1 - ) - - # The compatibility/fallback main and index branches, overlapped on the aux stream. - def _main_norm_rope(): - qkv = packed_qkv if packed_qkv is not None else self.qkv_proj(hidden_states) - fused_qkv = self._fused_qk_norm_rope( - qkv, - position_ids, - num_heads_q=self.num_heads, - num_heads_k=self.num_key_value_heads, - num_heads_v=self.num_key_value_heads, - head_dim=self.head_dim, - q_norm=self.q_norm, - k_norm=self.k_norm, - out_fp8=self._emit_fp8_main_qkv(), - ) - if fused_qkv is not None: - return self._split_main_qkv(fused_qkv) - assert not self._expect_fused_qk_norm_rope(position_ids), ( - f"MiniMax-M3 sparse attention (layer {self.layer_idx}) expected the fused QK-norm+RoPE kernel" - ) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = self.apply_qk_norm(q, k) - if self.rotary_emb is not None and position_ids is not None: - q, k = self.rotary_emb(position_ids, [q, k]) - return q, k, v - - def _index_norm_rope(): - idx_qk = ( - packed_idx_qk if packed_idx_qk is not None else self.index_qk_proj(hidden_states) - ) - graph_fp8_indexer = ( - self.register_to_config - and is_torch_compiling() - and isinstance(self.attn, MiniMaxM3MsaSparseAttention) - and self.attn.indexer_kv_dtype == "fp8" - ) - if not graph_fp8_indexer: - fp8_idx_q = self._fused_fp8_index_qk_norm_rope(idx_qk, position_ids, attn_metadata) - if fp8_idx_q is not None: - return fp8_idx_q, None - fused_idx = self._fused_qk_norm_rope( - idx_qk, - position_ids, - num_heads_q=self.sparse_num_index_heads, - num_heads_k=1, - num_heads_v=0, - head_dim=self.sparse_index_dim, - q_norm=self.index_q_norm, - k_norm=self.index_k_norm, - out_fp8=graph_fp8_indexer, - ) - if fused_idx is not None: - return self._split_index_qk(fused_idx) - assert not self._expect_fused_qk_norm_rope(position_ids), ( - f"MiniMax-M3 sparse index branch (layer {self.layer_idx}) expected the fused QK-norm+RoPE kernel" - ) - idx_q, idx_k = idx_qk.split([self.index_q_size, self.index_k_size], dim=-1) - idx_q, idx_k = self.apply_index_qk_norm(idx_q, idx_k) - if self.rotary_emb is not None and position_ids is not None: - idx_q, idx_k = self.rotary_emb(position_ids, [idx_q, idx_k]) - return idx_q, idx_k - - (q, k, v), (idx_q, idx_k) = maybe_execute_in_parallel( - _main_norm_rope, - _index_norm_rope, - self.ln_events[0], - self.ln_events[1], - self.aux_stream, - disable_on_compile=True, - ) - o = self._forward_attention_core(q, k, v, idx_q, idx_k, attn_metadata) - return self.o_proj(o, all_reduce_params=all_reduce_params) - - -def _plain_capture(spec_metadata, layer_id: int, hidden_states, residual) -> None: - """The in-tree Eagle3 hidden-state capture (no fused fc_norm, no early fc fork).""" - capture_hidden_states(spec_metadata, layer_id, hidden_states, residual, None) - - -class M3DecoderLayer(DecoderLayer): - """One MiniMax-M3 transformer block on this target's modules. - - Mirrors the built-in ``MiniMaxM3DecoderLayer``; the layer boundaries add the - multicast all-reduce + residual + RMSNorm up to its token limits. - """ - - def __init__( - self, - model_config: ModelConfig[PretrainedConfig], - layer_idx: int, - aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - ) -> None: - super().__init__() - config = model_config.pretrained_config - self.hidden_size = config.hidden_size - self.layer_idx = layer_idx - self.mapping = model_config.mapping - self.enable_attention_dp = self.mapping.enable_attention_dp - - _, sparse_layer_ids = get_sparse_layer_ids(config) - disable_index_value_ids = set(get_sparse_disable_index_value_layer_ids(config)) - self.self_attn = M3Attention( - model_config=model_config, - layer_idx=layer_idx, - is_sparse_attention_layer=layer_idx in sparse_layer_ids, - disable_index_value=layer_idx in disable_index_value_ids, - aux_stream=aux_stream_dict[AuxStreamType.Attention], - ) - - _, moe_layer_ids = get_moe_layer_ids(config) - if layer_idx in moe_layer_ids: - self.block_sparse_moe = M3MoE( - model_config=model_config, aux_stream_dict=aux_stream_dict, layer_idx=layer_idx - ) - self.mlp = None - else: - dense_intermediate = int( - getattr(config, "dense_intermediate_size", config.intermediate_size) - ) - self.mlp = _dense_mlp(model_config, dense_intermediate, layer_idx) - self.block_sparse_moe = None - - # Plain (non-Gemma) boundary norms, so they can drive the fused all-reduce - # + residual + RMSNorm epilogue; the loader folds Gemma's (1 + weight) in. - self.input_layernorm = RMSNorm( - hidden_size=config.hidden_size, - eps=config.rms_norm_eps, - dtype=config.torch_dtype, - use_gemma=False, - ) - self.post_attention_layernorm = RMSNorm( - hidden_size=config.hidden_size, - eps=config.rms_norm_eps, - dtype=config.torch_dtype, - use_gemma=False, - ) - - # tp4 without attention DP: every boundary has a cross-rank reduction to fold. - self.pre_feed_forward_fusion = True - self.post_feed_forward_fusion = True - self.allreduce = AllReduce( - mapping=model_config.mapping, - strategy=model_config.allreduce_strategy, - dtype=config.torch_dtype, - ) - # Wired after the weight load to the next layer's input_layernorm (or the final norm). - self.next_layer_layernorm: Optional[RMSNorm] = None - self.next_is_decoder_layer = False - # The Eagle3 hidden-state capture of this layer; the target shell points it at its worker's (fused fc_norm, - # early fc fork). - self.capture: Callable[..., None] = _plain_capture - - def forward( - self, - position_ids: torch.IntTensor, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - residual: Optional[torch.Tensor], - spec_metadata: Optional[SpecMetadata] = None, - lora_params: Optional[dict] = None, - **kwargs, - ) -> Tuple[torch.Tensor, torch.Tensor]: - if residual is None: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - hidden_states = self.self_attn( - position_ids=position_ids, - hidden_states=hidden_states, - attn_metadata=attn_metadata, - all_reduce_params=AllReduceParams(enable_allreduce=False), - lora_params=lora_params, - **kwargs, - ) - capture = spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx) - if self.block_sparse_moe is not None: - return self.forward_moe( - hidden_states, attn_metadata, residual, capture, spec_metadata, lora_params - ) - return self.forward_mlp(hidden_states, residual, capture, spec_metadata, lora_params) - - def _can_mc(self, hidden_states: torch.Tensor, folded: bool = False) -> bool: - """Whether the multicast all-reduce + residual + RMSNorm takes this boundary.""" - return ( - isinstance(hidden_states, torch.Tensor) - and hidden_states.dim() == 2 - and hidden_states.shape[0] - <= (m3_mc_allreduce.MAX_TOKENS if folded else MC_AR_MAX_TOKENS) - and hidden_states.dtype == torch.bfloat16 - and hidden_states.is_contiguous() - ) - - def mc_allreduce_norm(self, x, x2, residual, norm): - """The multicast all-reduce + optional second addend + residual + RMSNorm, or None without a workspace.""" - workspace = m3_mc_allreduce.get_workspace(self.mapping, x.shape[-1]) - if workspace is None: - return None - return m3_mc_allreduce.m3_mc_allreduce_norm( - x, - residual.contiguous(), - norm.weight, - norm.variance_epsilon, - workspace, - use_pdl=get_env_enable_pdl(), - x2=x2, - ) - - def native_allreduce_norm(self, hidden_states, residual, norm): - """The built-in fused all-reduce + residual + RMSNorm.""" - return self.allreduce( - hidden_states, - all_reduce_params=AllReduceParams( - fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, - residual=residual, - norm_weight=norm.weight, - eps=norm.variance_epsilon, - trigger_completion_at_end=False, - ), - ) - - def _pre_feed_forward_norm(self, hidden_states, residual): - if self._can_mc(hidden_states): - fused = self.mc_allreduce_norm( - hidden_states, None, residual, self.post_attention_layernorm - ) - if fused is not None: - return fused - return self.native_allreduce_norm(hidden_states, residual, self.post_attention_layernorm) - - def _next_layer_norm(self, hidden_states, residual, fused_boundary: bool): - """The next layer's input norm at the boundary; ``fused_boundary`` False keeps the module's own reduction.""" - if self.next_layer_layernorm is None: - return hidden_states, residual - parts = hidden_states if isinstance(hidden_states, tuple) else (hidden_states, None) - if not fused_boundary: - return self.next_layer_layernorm(parts[0], residual) - if self._can_mc(parts[0], folded=parts[1] is not None): - fused = self.mc_allreduce_norm(parts[0], parts[1], residual, self.next_layer_layernorm) - if fused is not None: - return fused - hidden_states = parts[0] if parts[1] is None else parts[0].add_(parts[1]) - return self.native_allreduce_norm(hidden_states, residual, self.next_layer_layernorm) - - def _fused_feed_forward_boundary(self, capture: bool) -> bool: - """Whether the feed-forward output goes unreduced to the fused boundary (all-reduce + residual + next norm in - one kernel). An Eagle3 capture layer takes it too and captures the boundary's reduced residual sum; only the - last layer (no next norm, its raw output captured) keeps the module's own all-reduce.""" - return not capture or self.next_layer_layernorm is not None - - def _feed_forward_boundary( - self, hidden_states, residual, fused_boundary, capture, spec_metadata - ): - """The layer's exit: the next norm at the boundary, with this layer's Eagle3 capture.""" - if not capture: - return self._next_layer_norm(hidden_states, residual, fused_boundary) - if not fused_boundary: - self.capture(spec_metadata, self.layer_idx, hidden_states, residual) - return self._next_layer_norm(hidden_states, residual, False) - normed, updated = self._next_layer_norm(hidden_states, residual, True) - self.capture(spec_metadata, self.layer_idx, updated, None) - return normed, updated - - def forward_moe( - self, hidden_states, attn_metadata, residual, capture, spec_metadata, lora_params - ): - hidden_states, residual = self._pre_feed_forward_norm(hidden_states, residual) - fused_boundary = self._fused_feed_forward_boundary(capture) - hidden_states = self.block_sparse_moe( - hidden_states, - attn_metadata, - final_all_reduce_params=AllReduceParams(enable_allreduce=False) - if fused_boundary - else None, - lora_params=lora_params, - return_parts=fused_boundary and self._can_mc(hidden_states, folded=True), - ) - return self._feed_forward_boundary( - hidden_states, residual, fused_boundary, capture, spec_metadata - ) - - def forward_mlp(self, hidden_states, residual, capture, spec_metadata, lora_params): - hidden_states, residual = self._pre_feed_forward_norm(hidden_states, residual) - fused_boundary = self._fused_feed_forward_boundary(capture) - hidden_states = self.mlp( - hidden_states, - final_all_reduce_params=AllReduceParams(enable_allreduce=False) - if fused_boundary - else None, - lora_params=lora_params, - ) - return self._feed_forward_boundary( - hidden_states, residual, fused_boundary, capture, spec_metadata - ) - - -class M3TextModel(DecoderModel): - """The MiniMax-M3 text decoder on this target's layers. - - Its forward belongs to the core in ``modeling.py``, which dispatches each - step to the prefill or the decode target. - """ - - def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: - _validate_sparse_attention_runtime_config(model_config) - super().__init__(model_config) - config = model_config.pretrained_config - self.vocab_size = config.vocab_size - self.aux_stream_dict = { - AuxStreamType.Attention: torch.cuda.Stream(), - AuxStreamType.MoeShared: torch.cuda.Stream(), - AuxStreamType.MoeChunkingOverlap: torch.cuda.Stream(), - } - self.embed_tokens = Embedding( - config.vocab_size, - config.hidden_size, - dtype=config.torch_dtype, - enable_torch_compile_for_embedding=model_config.enable_torch_compile_for_embedding, - ) - self.layers = nn.ModuleList( - [ - M3DecoderLayer(model_config, layer_idx, self.aux_stream_dict) - for layer_idx in range(config.num_hidden_layers) - ] - ) - # Lets the executor enable the automatic FlashInfer MXFP8 path for the decode graphs. - self._use_flashinfer_mxfp8_decode_graph_default = True - self.norm = RMSNorm( - hidden_size=config.hidden_size, - eps=config.rms_norm_eps, - dtype=config.torch_dtype, - use_gemma=False, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/linears.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/linears.py deleted file mode 100644 index a36255fa92e2..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/linears.py +++ /dev/null @@ -1,259 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The decode-graph MXFP8 GEMM dispatch of this target's projections. - -Every MXFP8 projection of the text model (QKV, o_proj, the dense MLPs and the -shared experts) runs the built-in ``MXFP8LinearMethod`` outside decode CUDA -graphs. Inside them (``flashinfer_mxfp8_decode_graph_capture_active``) the -projection's (N, K) and the step's rows pick one path (``dispatch.linear_path``, -row ranges per projection in ``dispatch.py``): - -1. an activation its producer already quantized to MXFP8 (the boundary - all-reduce or the attention merge publish it in ``m3_prequant.SLOT``): the - tcgen05 decode GEMV on the rows it serves, else the decode-graph MXFP8 GEMM - (AutoTuner-selected backend per token bucket, the built-in method's GEMM); -2. a bf16 activation on the rows the GEMV serves quantizing in its prologue; -3. projections marked for it, on its rows: the fused low-M MXFP8 GEMM; -4. otherwise the built-in method (MXFP8 quantize + GEMM, AutoTuner-selected per bucket). - -No path runs FlashInfer's own per-rank config sweep: its tactics are not merged -across ranks and it profiled wherever the projection ran (the MoE aux stream for -the shared expert), one CUDA graph per tactic. - -A kernel declining rows its path covers is a miss (``dispatch.missed``); the -chain then continues down the list. The GEMVs push their rows straight into -the boundary all-reduce the decode forward armed for the projection's push -slot (``m3_push``, fixed at construction). -""" - -from __future__ import annotations - -from typing import List, Optional, Tuple - -import torch -from torch import nn - -from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_gemm_autotuned -from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.mx_gemv import op as mx_gemv -from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl -from tensorrt_llm._torch.modules.linear import ( - Linear, - MXFP8LinearMethod, - flashinfer_mxfp8_autotune_active, - flashinfer_mxfp8_decode_graph_capture_active, -) -from tensorrt_llm._torch.utils import is_torch_compiling - -from . import dispatch -from .kernels import m3_mc_allreduce, m3_prequant -from .kernels.mxfp8_low_m_gemm import mxfp8_low_m_gemm -from .ops.gemm.m3_mx_gemv import m3_mx_gemv -from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 - -# The boundary all-reduce a projection's GEMV pushes its rows into, when the decode forward arms one: the attention -# boundary (o_proj) or the MoE boundary's shared-expert set (the shared expert's down projection). -PUSH_ATTN = "attn" -PUSH_MOE = "moe" - - -def mark_m3_linear(module: Linear, low_m_gemm: bool, push: Optional[str] = None) -> None: - """Mark a projection for this target's decode dispatch; ``adopt_m3_linears`` applies it after the load. - - ``low_m_gemm``: the projection may take the fused low-M MXFP8 GEMM. ``push``: ``PUSH_ATTN`` / ``PUSH_MOE`` when - its GEMV pushes into that boundary's all-reduce whenever the decode forward armed it. - """ - module.m3_low_m_gemm = low_m_gemm - module.m3_push = push - - -def adopt_m3_linears(root: nn.Module) -> None: - """Put every marked projection under ``root`` on the decode dispatch. - - Runs once the weights are real: the quantization method a projection ends - up with is only final after the model applied its quantization config. - It also compiles every decode GEMV kernel a projection can reach here, - before the executor's warmup: on the current (main) stream, the same shapes - in the same order on every rank, and never inside a forward between - collectives or on the MoE aux stream (``prebuild_m3_gemvs``). - """ - use_pdl = get_env_enable_pdl() - compiled = False - for name, module in root.named_modules(): - if not hasattr(module, "m3_low_m_gemm"): - continue - method = module.quant_method - if not isinstance(method, MXFP8LinearMethod) or not method.use_cutlass: - raise ValueError( - f"{name}: the MiniMax-M3 decode dispatch needs a CUTLASS MXFP8 projection, got {type(method).__name__}" - ) - if type(method) is MXFP8LinearMethod: - method.__class__ = M3MXFP8LinearMethod - mx_gemv.ensure_compiled(module.weight, module.weight_scale, use_pdl) - prebuild_m3_gemvs(module, use_pdl) - compiled = True - if compiled: - torch.cuda.synchronize() - - -def _bucket_tokens(rows: int) -> int: - """The token capacity of the multicast all-reduce bucket a ``rows``-token boundary uses.""" - return next(b for b in m3_mc_allreduce.BUCKETS if rows <= b) - - -def gemv_variants(module: Linear) -> List[Tuple[int, bool, Optional[Tuple[int, int, int]]]]: - """Every (rows, bf16 activation, push (world, bucket tokens, mode)) call the decode dispatch can route to this - projection's GEMV (``M3MXFP8LinearMethod.apply``): 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.""" - n, k = module.weight.shape - push_rows = {PUSH_ATTN: dispatch.ATTN_PUSH.keys, PUSH_MOE: dispatch.MOE_PUSH.keys}.get( - module.m3_push, frozenset() - ) - world = int(module.tp_size) - variants = [] - for rows in range(1, dispatch.DECODE_MAX_TOKENS + 1): - for x_bf16, served in ( - (False, rows in dispatch.MX_GEMV_ROWS.get((n, k), ())), - (True, rows <= dispatch.MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0)), - ): - if not served: - continue - variants.append((rows, x_bf16, None)) - if rows in push_rows: - variants.append( - (rows, x_bf16, (world, _bucket_tokens(rows), dispatch.AR_PUSH_MODE)) - ) - return variants - - -def prebuild_m3_gemvs(module: Linear, use_pdl: bool) -> int: - """Compile every GEMV kernel of ``gemv_variants`` (the push ones included) without launching; returns how many - distinct kernels the projection takes.""" - keys = set() - for rows, x_bf16, push in gemv_variants(module): - key = mx_gemv.precompile( - module.weight, module.weight_scale, rows, x_bf16, use_pdl=use_pdl, ar=push - ) - if key is not None: - keys.add(key) - return len(keys) - - -def _low_m_gemm_supports(x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor) -> bool: - """Whether the fused MXFP8 low-M kernel handles these operands.""" - return ( - x.dim() == 2 - and x.dtype in (torch.bfloat16, torch.float16) - and x.stride(-1) == 1 - and weight.dtype == torch.float8_e4m3fn - and weight.is_contiguous() - and weight_scale.dtype == torch.uint8 - and x.shape[1] == weight.shape[1] - and weight.shape[1] % 256 == 0 - ) - - -def _push_target(module: Linear): - """The armed boundary all-reduce this projection pushes its output into, if any.""" - if module.m3_push == PUSH_MOE: - return m3_prequant.MOE_PUSH.take_shared() - if module.m3_push == PUSH_ATTN: - return m3_prequant.AR_PUSH.take_target() - return None - - -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 _note_pushed(output: Optional[torch.Tensor], module: Linear) -> None: - if output is None or not mx_gemv.last_pushed: - return - if module.m3_push == PUSH_MOE: - m3_prequant.MOE_PUSH.shared_pushed = output.data_ptr() - else: - m3_prequant.AR_PUSH.pushed = output.data_ptr() - - -class M3MXFP8LinearMethod(MXFP8LinearMethod): - """``MXFP8LinearMethod`` with the decode-graph dispatch in the module docstring.""" - - def apply(self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor]): - if bias is not None or input.dim() != 2: - return super().apply(module, input, bias) - input = input.contiguous() - weight, weight_scale = module.weight, module.weight_scale - use_pdl = get_env_enable_pdl() - decode_graph = not is_torch_compiling() and flashinfer_mxfp8_decode_graph_capture_active() - if not decode_graph: - return super().apply(module, input, bias) - - rows = input.shape[0] - n, k = weight.shape - prequant = m3_prequant.SLOT.take(input) - if module.dtype != torch.bfloat16 or (prequant is None and input.dtype != torch.bfloat16): - path = ( - dispatch.PREQUANT_GEMM_PATH if prequant is not None else dispatch.GENERIC_GEMM_PATH - ) - else: - path = dispatch.linear_path(n, k, rows, prequant is not None, module.m3_low_m_gemm) - output = None - if prequant is not None: - act, act_scale = prequant - if path == dispatch.MX_GEMV_PATH: - output = m3_mx_gemv( - act, act_scale, weight, weight_scale, use_pdl=use_pdl, ar=_push_target(module) - ) - _note_pushed(output, module) - if output is None: - dispatch.missed( - path, - rows, - f"the GEMV declined the [{n}, {k}] projection's MXFP8 activation" - + _uncompiled_note(), - ) - if output is None: - if path == dispatch.PREQUANT_GEMM_PATH: - dispatch.uncovered( - dispatch.MX_GEMV_PATH, - f"the [{n}, {k}] projection at rows the GEMV does not serve", - ) - output = mxfp8_gemm_autotuned( - act, - act_scale, - weight, - weight_scale, - module.dtype, - tune=flashinfer_mxfp8_autotune_active(), - ) - return output - if path == dispatch.MX_GEMV_BF16_PATH: - output = m3_mx_gemv_bf16( - input, weight, weight_scale, use_pdl=use_pdl, ar=_push_target(module) - ) - _note_pushed(output, module) - if output is None: - dispatch.missed( - path, - rows, - f"the GEMV declined the [{n}, {k}] projection's input" + _uncompiled_note(), - ) - if ( - output is None - and path in (dispatch.MX_GEMV_BF16_PATH, dispatch.LOW_M_GEMM_PATH) - and module.m3_low_m_gemm - and rows <= dispatch.LOW_M_GEMM_MAX_TOKENS.get((n, k), 0) - and _low_m_gemm_supports(input, weight, weight_scale) - ): - output = mxfp8_low_m_gemm(input, weight, weight_scale, module.dtype, use_pdl=use_pdl) - if output is None and path == dispatch.LOW_M_GEMM_PATH: - dispatch.missed(path, rows, f"the GEMM declined the [{n}, {k}] projection's input") - if output is None: - if path == dispatch.GENERIC_GEMM_PATH: - dispatch.uncovered( - dispatch.MX_GEMV_BF16_PATH, - f"the [{n}, {k}] projection's bf16 input beyond the decode GEMV's rows", - ) - # Stock's decode-graph dispatch for this shape (mxfp8_quantize_gemm_autotuned): same tactics, same bits. - output = super().apply(module, input, bias) - return output diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py index 219d0a1e0300..e627ac7a7419 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py @@ -1,568 +1,4530 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""ModelingV2 target: MiniMax-M3 NVFP4 / sm_103 / tp4. - -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) -plus one shared expert. The checkpoint is a modelopt mixed-precision export: -NVFP4 routed experts, MXFP8 projections, bf16 router, embedding and head; the KV -cache is NVFP4 for the sparse layers and FP8 for the dense ones. - -``tp4`` is four GB300s, head-split attention (16 query heads and 1 KV head per -rank) and tensor-parallel experts (every expert on every rank, width 768 per -rank), no attention DP. Speculative decoding, when configured, is one-model -Eagle3 with draft length 3: the in-tree BF16 drafter over the full vocabulary -and the BF16 target head, both at the checkpoint's precision; decode steps of -up to 8 requests verify and draft through fused kernels at that precision -(``spec_worker.py``). - -**Each step is dispatched once, on the host, from its shape** (``_select_target``): - -* ``DecodeTarget`` takes a pure generation step of at most 48 tokens (up to 8 - requests, each one token or a draft verify of 1 + 3 + extra rows). It runs - the layers with the decode-step kernels of ``layers.py`` and chooses each - layer boundary's all-reduce by the step's token bucket (``DECODE_PLANS``). - Every fast path it can take, the rows each covers and what a step that cannot - take one does (raise under ``modeling_v2='require'``, else log once and fall - back) are in ``dispatch.py``. -* ``PrefillTarget`` takes every other step (prefill, mixed steps, wider decode - steps) on the built-in modules' generic path. - -**What this target asserts rather than adapts**: SM 10.3, the topology above, -the checkpoint shape and quantization (``routing.py``), the MiniMax-M3 MSA -sparse attention backend with its NVFP4 hybrid KV cache, and text-only input: -the vision tower is not built and a step carrying image input raises. +"""ModelingV2 target: MiniMax-M3 NVFP4 / sm_103 / 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 GB300s: 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 fused +decode attention for the o_proj. It publishes it in ``kernels/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 typing import Any, Dict, List, Optional +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.cute_dsl_kernels.minimax_m3.mx_gemv import op as mx_gemv -from tensorrt_llm._torch.distributed import AllReduceParams +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.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_minimaxm3 import ( - MiniMaxM3ForCausalLM, - get_text_model_config, - is_minimax_m3_vl_config, +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.models.modeling_speculative import SpecDecOneEngineForCausalLM -from tensorrt_llm._torch.modules.linear import 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.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.functional import PositionEmbeddingType from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantAlgo -from . import dispatch from . import weights as _weights -from .dispatch import AR_PUSH_MODE, EMBED_FIRSTNORM_MAX_TOKENS, DecodePlan, plan_for -from .kernels import m3_mc_allreduce, m3_prequant -from .layers import M3DecoderLayer, M3MoE, M3TextModel, prebuild_index_scorer -from .linears import adopt_m3_linears -from .ops.moe.m3_moe16_engine import m3_moe16_engine +from .kernels import m3_attn_mega as _attn_mega_kernel +from .kernels import m3_attn_mega_c1 as _attn_c1_kernel +from .kernels import m3_draft_step as m3_draft_step_kernel +from .kernels import m3_fused_attn as _fused_attn_kernel +from .kernels import m3_index_score as _index_score_kernel +from .kernels import m3_mc_allreduce, m3_prequant, m3_spec_tail +from .kernels import m3_sampled_accept as m3_sampled_accept_kernel +from .kernels import m3_shard_accept as m3_shard_accept_kernel +from .kernels import m3_target_head as m3_target_head_kernel +from .kernels import msa_kv_update as _msa_kv_update_kernel +from .kernels.long_select import long_select_blocks +from .kernels.m3_router_gemv import m3_router_gemv as _router_gemv +from .kernels.m3_router_gemv import m3_router_gemv_fp4 as _router_gemv_fp4 +from .kernels.m3_router_gemv import m3_router_tc as _router_tc +from .kernels.mx_gemv import op as mx_gemv +from .kernels.mxfp8_low_m_gemm import mxfp8_low_m_gemm as _low_m_gemm +from .kernels.triton_sparse_decode import MERGE_MXFP8, MERGE_MXFP8_MIN_TOKENS +from .ops.activation.m3_silu_and_mul import m3_silu_and_mul +from .ops.attention.m3_attn_c1 import m3_attn_c1 +from .ops.attention.m3_attn_mega import m3_attn_mega +from .ops.attention.m3_dense_decode import m3_dense_decode +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_fused_decode_attn import m3_fused_decode_attn +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_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 MAX_BLOCKS as SELECT_LONG_MAX_BLOCKS +from .ops.attention.m3_select_long import m3_select_long +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_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_target_head import MAX_ROWS as TARGET_HEAD_KERNEL_MAX_ROWS +from .ops.gemm.m3_target_head import m3_target_head +from .ops.gemm.mxfp8_low_m_gemm import mxfp8_low_m_gemm +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_gemv import m3_router_gemv +from .ops.moe.m3_router_gemv_fp4 import m3_router_gemv_fp4 +from .ops.moe.m3_router_tc import m3_router_tc from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm -from .runtime import make_runtime_hooks -from .spec_worker import M3Eagle3Worker, M3SpecDecForwardMixin +from .ops.spec.m3_draft_input import m3_draft_input +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_shard_accept import m3_shard_accept _SM = (10, 3) +# -------------------------------------------------------------------------------------------------------------------- +# Row ranges +# -------------------------------------------------------------------------------------------------------------------- +# The largest generation step the decode kernels take: 8 requests verifying 1 + 3 drafts and their padding rows. A +# wider pure generation step stays on the decode target and runs its generic layers. +DECODE_MAX_TOKENS = 48 # capacity: fixed: the decode kernels' step bound -class MiniMaxM3Tp4ModelingV2(M3TextModel): - """The text model of this target: owns the layers and dispatches each step to a target.""" - def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: - super().__init__(model_config) - self._embed_owners: Dict[tuple, DecodeEmbedding] = {} - self._embed_identity: Optional[tuple] = None - self._prefill: Optional[PrefillTarget] = None - self._decode: Optional[DecodeTarget] = None +@dataclass(frozen=True) +class DecodePlan: + """The layer-boundary collectives of one decode token bucket (GB300 tp4 measurements in the comments).""" - def build_targets(self) -> None: - """Build the two targets, once the weights are real.""" - table = self.embed_tokens.weight - norm = self.layers[0].input_layernorm - identity = ( - id(table), - table.data_ptr(), - id(norm.weight), - norm.weight.data_ptr(), - float(norm.variance_epsilon), + 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 (kernels/m3_mc_allreduce.BUCKETS). +DECODE_PLANS = ( + DecodePlan("tiny", 8, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), + DecodePlan("small", 16, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), + DecodePlan("verify", DECODE_MAX_TOKENS, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), +) + +# 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. +DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({1: 16384}) + +# 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 + +# 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 = 16 # capacity: fixed: measured push crossover +AR_PUSH_MODE = 2 +# The 1..8-token MoE engine (router, routed experts and the shared expert in one launch) pushes its sum into the MoE +# boundary's all-reduce (prepushed mode 2). +MOE_ENGINE_MAX_TOKENS = 8 # capacity: fixed: the engine's built rows +MOE_PUSH_MAX_TOKENS = MOE_ENGINE_MAX_TOKENS # 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). +MC_FINALIZE_MIN_TOKENS = 9 # capacity: fixed: the rows above the engines +# The router: the fp32 GEMV at one row, the tensor-core router (fused with the routed experts' NVFP4 activation +# quantization) at 2..48 rows. GB300 tp4, behind a PDL producer, with the NVFP4 activation: 2.90 / 2.93 / 3.20 / +# 3.82 / 7.36 us at 1 / 4 / 8 / 12 / 48 tokens against the fp32 GEMV's 2.85 / 3.96 / 4.88 / 6.01 / - us. +ROUTER_GEMV_MAX_TOKENS = 1 # capacity: fixed: measured router crossover +ROUTER_TC_MAX_TOKENS = 48 # capacity: fixed: the tensor-core router's built rows +# The fused norm/RoPE/KV-insert/sparse-decode/merge kernel: up to 8 query rows per request and 16 query rows per +# step (above, its 16-CTA clusters run in waves and the three-kernel path is faster: GB300 -13% fAL spec c8). +FUSED_ATTN_MAX_QUERY_LEN = 8 # capacity: fixed: compiled query rows +FUSED_ATTN_MAX_TOKENS = 16 # capacity: fixed: measured wave crossover +# One request verifying a draft of 6 query rows runs attention + o_proj as one launch (96 CTAs: 6 clusters of 16). +ATTN_C1_QUERY_LEN = 6 # capacity: fixed: the verify kernel's query rows +# 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 (kernels/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 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 (kernels/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. +TARGET_HEAD_MAX_ROWS = TARGET_HEAD_KERNEL_MAX_ROWS +# 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 = 8 # capacity: fixed: measured fork crossover +# 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 (kernels/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 (kernels/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 = True +DRAFT_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN # capacity: fixed: the built draft length +DRAFT_STEP_MAX_REQUESTS = DRAFT_STEP_KERNEL_MAX_TOKENS +# The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows +# (kernels/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 = SELECT_LONG_MAX_BLOCKS +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; on wider steps the early-launched attention CTAs crowd +# the select's CTAs off the SMs (8 requests x 4 rows: +3.1 us per layer), so they keep the trigger at the select's exit. +SELECT_EARLY_MAX_ROWS = 8 # capacity: fixed: measured early-release crossover +# 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 (larger batches +# are bandwidth bound, where the extra early tiles only crowd the prefetch). +INDEX_SCORE_EARLY_MAX_REQUESTS = 1 # capacity: fixed: measured early-issue crossover + +# 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. +_ALL_ROWS = frozenset(range(1, DECODE_MAX_TOKENS + 1)) +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, +} +# 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, +} +# The fused low-M MXFP8 GEMM (projections marked for it), per projection: its built configurations' row limit +# (kernels/mxfp8_low_m_gemm). The bf16 GEMV serves every one of these rows first; the GEMM takes them only if it +# declines. +LOW_M_GEMM_MAX_TOKENS: Dict[Tuple[int, int], int] = { + QKV_SPARSE: 8, + O_PROJ: 8, + DENSE_GATE_UP: 1, + DENSE_DOWN: 8, +} + +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 _fused_attn_keys() -> FrozenSet[Tuple[int, int]]: + return frozenset( + (qlen, reqs) + for qlen in range(1, FUSED_ATTN_MAX_QUERY_LEN + 1) + for reqs in range(1, FUSED_ATTN_MAX_TOKENS // qlen + 1) + ) + + +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"), +) +# One request decoding one token on an MXFP8 layer input: QKV GEMV + attention + o_proj in one launch. +ATTN_MEGA = RowRange("attn_mega", frozenset({(1, 1)}), m3_attn_mega, _const((1, 1))) +ATTN_C1 = RowRange( + "attn_c1", frozenset({(ATTN_C1_QUERY_LEN, 1)}), m3_attn_c1, _const((ATTN_C1_QUERY_LEN, 1)) +) +FUSED_DECODE_ATTN = RowRange( + "fused_decode_attn", + _fused_attn_keys(), + m3_fused_decode_attn, + lambda spec: ((spec["qlen"], len(spec["seq_lens"])),), +) +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",), +) +LONG_SELECT = RowRange( + "long_select", frozenset(LONG_SELECT_KEY_WIDTHS), m3_select_long, _long_select_widths +) +MC_ALLREDUCE = RowRange( + "mc_allreduce_norm", + _rows(1, DECODE_MAX_TOKENS), + 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), +) +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, DECODE_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"), +) +ROUTER_GEMV = RowRange( + "m3_router_gemv", _rows(1, ROUTER_GEMV_MAX_TOKENS), m3_router_gemv, _spec_rows("m") +) +ROUTER_GEMV_FP4 = RowRange( + "m3_router_gemv_fp4", _rows(1, ROUTER_GEMV_MAX_TOKENS), m3_router_gemv_fp4, _spec_rows("m") +) +ROUTER_TC = RowRange( + "m3_router_tc", + _rows(ROUTER_GEMV_MAX_TOKENS + 1, ROUTER_TC_MAX_TOKENS), + m3_router_tc, + _spec_rows("m"), +) +# 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"]),), +) +LOW_M_GEMM = RowRange( + "mxfp8_low_m_gemm", + frozenset((nk, m) for nk, hi in LOW_M_GEMM_MAX_TOKENS.items() for m in _rows(1, hi)), + mxfp8_low_m_gemm, + lambda spec: (((spec["n"], spec["k"]), 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 (), +) + +# 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"]),), +) +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, + ATTN_MEGA, + ATTN_C1, + FUSED_DECODE_ATTN, + DENSE_DECODE_SPLIT, + SPARSE_DECODE_PRODUCER, + PRODUCER, + INDEX_SCORE, + LONG_SELECT, + MC_ALLREDUCE, + ATTN_PUSH, + MOE_PUSH, + MC_FINALIZE, + MOE_ENGINE, + MOE16_ENGINE, + ROUTER_GEMV, + ROUTER_GEMV_FP4, + ROUTER_TC, + MX_GEMV, + MX_GEMV_BF16, + LOW_M_GEMM, + SHARD_ACCEPT, + TARGET_HEAD, + DRAFT_STEP, + 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 + + +def router_path(num_tokens: int) -> Optional[str]: + """The router kernel when the MoE runs unfused: the fp32 GEMV, the tensor-core router, or None (built-in gate).""" + if ROUTER_GEMV.covers(num_tokens): + return ROUTER_GEMV.path + if ROUTER_TC.covers(num_tokens): + return ROUTER_TC.path + return None + + +MX_GEMV_PATH = MX_GEMV.path +MX_GEMV_BF16_PATH = MX_GEMV_BF16.path +LOW_M_GEMM_PATH = LOW_M_GEMM.path +# 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``: the projection is marked for the + fused low-M GEMM. + """ + 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 low_m_gemm and num_tokens <= LOW_M_GEMM_MAX_TOKENS.get((n, k), 0): + return LOW_M_GEMM_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}" ) - if self._embed_owners and identity != self._embed_identity: - raise RuntimeError( - "the embedding or the first norm moved after the decode embedding owners were built" - ) - if not self._embed_owners: - self._embed_owners = { - (rows, dtype): DecodeEmbedding(table, norm, rows, dtype) - for rows in range(1, EMBED_FIRSTNORM_MAX_TOKENS + 1) - for dtype in (torch.int32, torch.int64) - } - self._embed_identity = identity - self.prebuild_kernels() - self._prefill = PrefillTarget(self) - self._decode = DecodeTarget(self) + _log_once( + logger.warning, + path, + reason, + f"[M3 dispatch] {path}{at}: {reason}; falling back (logged once per reason)", + ) - 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 ``dispatch.missed``). - Returns the seconds each part took and the scorer keys.""" - times: Dict[str, object] = {} - mx_gemv.freeze(False) - start = time.perf_counter() - adopt_m3_linears(self) - times["gemv_s"] = time.perf_counter() - start - start = time.perf_counter() - for layer in self.layers: - moe = layer.block_sparse_moe - if isinstance(moe, M3MoE): - moe.prebuild_engines() - torch.cuda.synchronize() - times["moe_s"] = time.perf_counter() - start - start = time.perf_counter() - times["scorer_keys"] = prebuild_index_scorer( - self.layers, self.decode_query_lens(), self.embed_tokens.weight.device +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 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() + + +# -------------------------------------------------------------------------------------------------------------------- +# 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 (m3_mc_allreduce.MAX_TOKENS). +MC_AR_MAX_TOKENS = 8 # capacity: fixed: measured crossover vs the one-shot all-reduce +# 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) + +# -------------------------------------------------------------------------------------------------------------------- +# 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 not in the checkpoint): routing declines every other one + (``routing._runtime_mismatch``), so this is the backstop for a construction that did not go through routing. + + 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'" ) - torch.cuda.synchronize() - times["scorer_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)" - % (times["gemv_s"], times["moe_s"], times["scorer_s"], len(times["scorer_keys"])) + 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, "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 _low_m_gemm_supports(x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor) -> bool: + """Whether the fused MXFP8 low-M kernel handles these operands.""" + return ( + x.dim() == 2 + and x.dtype in (torch.bfloat16, torch.float16) + and x.stride(-1) == 1 + and weight.dtype == torch.float8_e4m3fn + and weight.is_contiguous() + and weight_scale.dtype == torch.uint8 + and x.shape[1] == weight.shape[1] + and weight.shape[1] % 256 == 0 + ) + + +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. ``low_m_gemm`` projections, on its rows: the fused low-M MXFP8 GEMM; + 4. 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. + """ + 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 + 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: + 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") + 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: + 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 in (MX_GEMV_BF16_PATH, LOW_M_GEMM_PATH) + and low_m_gemm + and rows <= LOW_M_GEMM_MAX_TOKENS.get((n, k), 0) + and _low_m_gemm_supports(x, w, s) + ): + output = _low_m_gemm(x, w, s, torch.bfloat16, use_pdl=use_pdl) + if output is not None: + return output, False + if path == LOW_M_GEMM_PATH: + missed(path, rows, f"the GEMM declined the [{n}, {k}] projection's input") + if path == GENERIC_GEMM_PATH: + uncovered( + MX_GEMV_BF16_PATH, + f"the [{n}, {k}] projection's bf16 input beyond the decode GEMV's rows", + ) + 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.""" + return ( + isinstance(x, torch.Tensor) + and x.dim() == 2 + and x.shape[0] <= (m3_mc_allreduce.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): + """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.""" + 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, + ) + + +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``: the fp32 GEMV at one row, the tensor-core router up to 48, else the + fp32 linear.""" + rows = h.shape[0] if h.dim() == 2 else 0 + path = router_path(rows) + w = core.w[f"l{i}_router"] + if path == ROUTER_GEMV.path: + logits = _router_gemv(h, w, use_pdl=get_env_enable_pdl()) + if logits is not None: + return logits + missed(path, rows, "the router GEMV declined the operands") + elif path == ROUTER_TC.path: + fused = _router_tc(h, w, use_pdl=get_env_enable_pdl()) + if fused is not None: + return fused[0] + missed(path, rows, "the tensor-core router declined the operands") + else: + uncovered(ROUTER_TC.path, f"more than {ROUTER_TC_MAX_TOKENS} rows") + return m3_fp32_linear(h, w) + + +def _router_fp4(core, i: int, h: torch.Tensor): + """Router logits and the routed experts' NVFP4 activation from one kernel, or None (then ``_router_logits``).""" + if h.dim() != 2 or h.shape[0] > ROUTER_TC_MAX_TOKENS: + return None + rows = h.shape[0] + path = router_path(rows) + if path == ROUTER_GEMV.path: + path = ROUTER_GEMV_FP4.path + scale = core.moe[i].fc31_input_scale + if path == ROUTER_TC.path: + fused = _router_tc(h, core.w[f"l{i}_router"], scale, use_pdl=get_env_enable_pdl()) + else: + fused = _router_gemv_fp4(h, core.w[f"l{i}_router"], scale, use_pdl=get_env_enable_pdl()) + if fused is None: + missed(path, rows, "the router declined the operands") + return fused + + +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). 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) + + +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 (core.heads, core.kv_heads, core.index_heads) == (16, 1, 1) + 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, + ) + 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 _attn_mega(core, i: int, hidden: torch.Tensor, args: dict, push: Optional[tuple]): + """(o_proj output, pushed) of the one-launch QKV + attention + o_proj kernel for one request decoding one token + whose layer input the preceding all-reduce handed over in MXFP8; None, with nothing consumed, otherwise.""" + if not ATTN_MEGA.covers((args["decode_query_len"], int(args["seq_lens"].shape[0]))): + return None + if not flashinfer_mxfp8_decode_graph_capture_active(): + return None + pending = m3_prequant.SLOT.pending + if ( + pending is None + or hidden.dim() != 2 + or pending[0] != hidden.data_ptr() + or pending[1] != 1 + or pending[2] != hidden.shape[1] + ): + # After a layer whose boundary does not quantize (a dense layer, an Eagle3 capture layer). + uncovered(ATTN_MEGA.path, "the layer input has no MXFP8 hand-off") + return None + w = core.w + if w[f"l{i}_qkv_w"].shape != (2560, 6144) or w[f"l{i}_o_w"].shape != (6144, 2048): + uncovered(ATTN_MEGA.path, "not the tp4 QKV / o_proj shapes") + return None + workspace = core.mega_ws.get(i) + if workspace is None: + # The first call JIT-builds the module; both only outside capture (the eager warmup steps). + if torch.cuda.is_current_stream_capturing(): + capture_missed(ATTN_MEGA.path, 1, "the attn_mega module and workspace") + return None + _attn_mega_kernel._load_module() + workspace = core.mega_ws[i] = _attn_mega_kernel.MegaWorkspace(hidden.device) + xq, xsf = m3_prequant.SLOT.take(hidden) + output = hidden.new_empty((1, core.heads * core.head_dim), dtype=torch.bfloat16) + y = hidden.new_empty((1, w[f"l{i}_o_w"].shape[0]), dtype=torch.bfloat16) + _record("m3_attn_mega", push is not None) + _attn_mega_kernel.m3_attn_mega( + xq, + xsf, + w[f"l{i}_qkv_w"], + w[f"l{i}_qkv_s"], + workspace, + args["kv_data"], + args["kv_scale"], + args["index_k_cache"], + args["out_cache_loc"], + args["position_ids"], + args["block_table"], + args["seq_lens"], + args["kv_quant_scale"], + args["k_global_scale"], + args["v_global_scale"], + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_k_norm"], + core.rotary_cos_sin, + w[f"l{i}_o_w"], + w[f"l{i}_o_s"], + eps=core.rms_eps, + sm_scale=core.sm_scale(i), + decode_query_len=1, + output=output, + y=y, + use_pdl=get_env_enable_pdl(), + ar=push, + ) + return y, push is not None + + +def _attn_c1(core, i: int, packed: torch.Tensor, args: dict, push: Optional[tuple]): + """(o_proj output, True) of the one-launch attention + o_proj kernel for one request verifying a draft, pushed + into the attention boundary's all-reduce; None, with nothing consumed, when the step or the layer does not + qualify.""" + qlen, num_tokens = args["decode_query_len"], args["num_tokens"] + if not ATTN_C1.covers((qlen, int(args["seq_lens"].shape[0]))): + return None + if not flashinfer_mxfp8_decode_graph_capture_active(): + return None + if push is None: + missed(ATTN_C1.path, num_tokens, "no armed attention-boundary push to write into") + return None + w = core.w + if w[f"l{i}_o_w"].shape != (6144, 2048): + uncovered(ATTN_C1.path, "not the tp4 o_proj shape") + return None + workspace = core.c1_ws.get(i) + if workspace is None: + # The first call JIT-builds the module; both only outside capture (the eager warmup steps). + if torch.cuda.is_current_stream_capturing(): + capture_missed(ATTN_C1.path, num_tokens, "the attn_c1 module and workspace") + return None + _attn_c1_kernel._load_module() + workspace = core.c1_ws[i] = _attn_c1_kernel.MegaWorkspace(packed.device) + output = packed.new_empty((num_tokens, core.heads * core.head_dim), dtype=torch.bfloat16) + y = packed.new_empty((num_tokens, w[f"l{i}_o_w"].shape[0]), dtype=torch.bfloat16) + _record("m3_attn_c1", num_tokens) + _attn_c1_kernel.m3_attn_c1( + packed.contiguous(), + workspace, + args["kv_data"], + args["kv_scale"], + args["index_k_cache"], + args["out_cache_loc"], + args["position_ids"], + args["block_table"], + args["seq_lens"], + args["kv_quant_scale"], + args["k_global_scale"], + args["v_global_scale"], + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_k_norm"], + core.rotary_cos_sin, + w[f"l{i}_o_w"], + w[f"l{i}_o_s"], + eps=core.rms_eps, + sm_scale=core.sm_scale(i), + decode_query_len=qlen, + output=output, + y=y, + use_pdl=get_env_enable_pdl(), + ar=push[:5], + ) + return y, True + + +def _fused_attn(core, i: int, packed: torch.Tensor, args: dict) -> torch.Tensor: + """The fused norm/RoPE/KV-insert/sparse-decode/merge kernel after the packed QKV projection; its MXFP8 copy of + the output goes to the o_proj through ``m3_prequant.SLOT``.""" + num_tokens = args["num_tokens"] + width = core.heads * core.head_dim + output = packed.new_empty((num_tokens, width), dtype=torch.bfloat16) + mxfp8 = None + if MERGE_MXFP8_MIN_TOKENS <= num_tokens <= MERGE_MXFP8.MAX_ROWS: + mxfp8 = MERGE_MXFP8.buffers(width, output.device) + w = core.w + _record("m3_fused_decode_attn", num_tokens, args["decode_query_len"]) + _fused_attn_kernel.m3_fused_decode_attn( + packed.contiguous(), + args["kv_data"], + args["kv_scale"], + args["index_k_cache"], + args["out_cache_loc"], + args["position_ids"], + args["block_table"], + args["seq_lens"], + args["kv_quant_scale"], + args["k_global_scale"], + args["v_global_scale"], + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_k_norm"], + core.rotary_cos_sin, + eps=core.rms_eps, + sm_scale=core.sm_scale(i), + decode_query_len=args["decode_query_len"], + output=output, + mxfp8_out=mxfp8, + use_pdl=get_env_enable_pdl(), + ) + if mxfp8 is not None: + m3_prequant.SLOT.publish(output, mxfp8[0][:num_tokens], mxfp8[1]) + return output + + +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] = {} + self.mega_ws: Dict[int, object] = {} + self.c1_ws: Dict[int, object] = {} + 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", None)) + 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). Returns the + seconds each part took and the scorer keys.""" + times: Dict[str, object] = {} + if not self.mxfp8.quant_method.use_cutlass: + raise ValueError("the MiniMax-M3 decode dispatch needs the CUTLASS MXFP8 GEMM") + 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(), self.w["embed"].device + ) + torch.cuda.synchronize() + times["scorer_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)" + % (times["gemv_s"], times["moe_s"], times["scorer_s"], len(times["scorer_keys"])) + ) + 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 + + # --------------------------------------------------------------------------------------------- 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, + **kwargs, + ) -> torch.Tensor: + 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 + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# 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 + ): + 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) -> torch.Tensor: + """A gated MXFP8 MLP (the dense MLP, 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) + 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 fused router with the NVFP4 activation up to 48 rows, else the + router logits and the experts quantizing their bf16 input (chunked).""" + core = self.core + fused = _router_fp4(core, i, h) + if fused is not None: + router_logits, x_fp4, x_sf = fused + return _routed_experts(core, i, x_fp4, x_sf, router_logits, True) + router_logits = _router_logits(core, i, h) + return self._routed_chunks(i, h, router_logits) + + 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) + 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] = {} + # 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, + metadata=md, + ) + + def forward( + self, attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, capture + ): + 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): + if ( + not kernels + or residual is None + or (spec_metadata is not None and spec_metadata.is_layer_capture(i)) + ): + 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) + 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): + """All-reduce + (second addend) + residual + RMSNorm: the multicast kernel, else the built-in all-reduce.""" + core = self.core + if use_mc: + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w, quant=quant) + if fused is not None: + 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): + fused = _mc_allreduce_norm( + self.core, x, x2, residual, norm_w, quant=quant, prepushed=prepushed + ) + if fused is None: + raise RuntimeError(f"layer {i}: rows were pushed into an all-reduce that did not run") + return fused + + def _layer(self, i: int, plan: DecodePlan, position_ids, hidden, residual, step): + 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: + 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: + out, _ = self._swiglu_mlp(hidden, f"l{i}_mlp", True) + return self._boundary(plan.attn_mc_ar, out, residual, norm_w) + + 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: + 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) + return self._boundary(plan.moe_mc_ar, out, residual, norm_w, quant=quant) + 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" + ) + 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: + return self._prepushed(i, shared, residual, norm_w, quant, 2) + else: + if engine_pushed and shared_pushed: + return self._prepushed(i, shared, residual, norm_w, quant, 3, x2=routed) + 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) + + def _swiglu_mlp( + self, h: torch.Tensor, prefix: str, low_m_gemm: bool, push: Optional[tuple] = None + ): + """A gated MXFP8 MLP (the dense MLP, 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) + 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 fused router with the NVFP4 activation up to 48 rows, else the router logits and + the experts quantizing their bf16 input.""" + core = self.core + fused = _router_fp4(core, i, h) + if fused is not None: + router_logits, x_fp4, x_sf = fused + return _routed_experts(core, i, x_fp4, x_sf, router_logits, do_finalize) + router_logits = _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) + 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) + 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 MSA backend, o_proj. + 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 + fused = _qk_norm_rope(core, i, qkv, position_ids, out_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, push) + if not core.fused_projection: + return self._unfused_sparse_attention(i, position_ids, hidden, step, push) + # Gathered before the QKV projection so nothing runs between it and the fused attention (PDL hand-off). + fused = self._fused_attn_args(i, position_ids, step) + if fused is not None: + mega = _attn_mega(core, i, hidden, fused, push) + if mega is not None: + return mega + packed, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + if fused is not None and ( + packed.shape[0] != fused["num_tokens"] or packed.shape[1] != 2560 + ): + missed( + FUSED_DECODE_ATTN.path, + fused["num_tokens"], + f"the packed QKV is {tuple(packed.shape)}, not [rows, 2560]", + ) + fused = None + if fused is not None: + c1 = _attn_c1(core, i, packed, fused, push) + if c1 is not None: + return c1 + o = _fused_attn(core, i, packed, fused) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + 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 _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 _fused_attn_args(self, i: int, position_ids, step) -> Optional[dict]: + """Arguments of the fused decode attention for this step, or None when it does not apply. + + The layer must be an NVFP4 sparse MSA layer and the step a pure generation step whose selection the metadata + staged as the identity (every request within 16 blocks of 128 tokens), with at most 8 query tokens per request. + """ + core = self.core + path = FUSED_DECODE_ATTN.path + if i not in core.sparse_layers or not core.main_kv_is_nvfp4: + uncovered(path, "not an NVFP4 sparse MSA layer with an FP8 indexer") + return None + if position_ids is None: + uncovered(path, "not a pure generation step with positions") + return None + if not step["msa_short_ok"]: + uncovered( + path, "a request beyond 16 KV blocks: the block selection is not the identity" + ) + 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 + num_reqs = int(step["num_seqs"]) + qlen = int(span.query_len) + if not FUSED_DECODE_ATTN.covers((qlen, num_reqs)): + uncovered( + path, + f"more than {FUSED_ATTN_MAX_QUERY_LEN} query rows per request or {FUSED_ATTN_MAX_TOKENS} per step", + ) + return None + manager = step["kv_cache_manager"] + if manager is not None and not manager.is_nvfp4_layer(i): + uncovered(path, "the layer's KV cache is not NVFP4") + return None + if _fused_attn_kernel._module is None: + # Building compiles every geometry (minutes): only outside capture (the eager warmup steps). + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_reqs * qlen, "the fused attention module") + return None + _fused_attn_kernel._load_module() + if manager is None: + missed(path, num_reqs * qlen, "no KV-cache manager") + return None + k_scale, v_scale = _aligned_nvfp4_dequant_scales(core.attn[i], core.kv_scales[i]) + return dict( + num_tokens=num_reqs * qlen, + decode_query_len=qlen, + kv_data=manager.get_buffers(i, kv_layout="HND").view(torch.uint8), + kv_scale=manager.get_block_scale_buffers(i, kv_layout="HND"), + index_k_cache=step["msa_idx_k_cache"](i), + out_cache_loc=step["msa_out_cache_loc"], + position_ids=position_ids.reshape(-1).contiguous().to(torch.int32), + block_table=step["msa_block_table"][:num_reqs], + seq_lens=step["msa_seq_lens_cuda"][:num_reqs], + kv_quant_scale=core.inv_kv_scales[i], + k_global_scale=k_scale, + v_global_scale=v_scale, + ) + + 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 + ) + 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, + ) + 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 + ) + 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=num_tokens <= SELECT_EARLY_MAX_ROWS, + ) + if table is None: + 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) + 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, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# Speculative decoding: one-model Eagle3 at checkpoint precision +# -------------------------------------------------------------------------------------------------------------------- + + +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 never stores the sampler's outputs itself, so + ``begin_step`` and ``sampler_rows_stored`` always report that nothing was staged. + """ + + def __init__(self) -> None: + self.store: Optional[object] = None + self.step_inputs: Optional[tuple] = 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) + + def bind_sampler_store(self, store: object) -> None: + self.store = store + + def begin_step(self, seq_slots: Optional[list], num_rows: int) -> bool: + return False + + 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: + return False + + 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) + + +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) + + @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 m3_bf16_linear(hs, self.w["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 = m3_bf16_linear(x, self.w["qkv"]) + _record("m3_draft_attention", int(md.num_tokens)) + o = m3_draft_attention(self.attn, qkv, md) + out = m3_bf16_linear(o, self.w["o"]) + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def mlp(self, h: torch.Tensor, reduce: bool) -> torch.Tensor: + gu = m3_bf16_linear(h, self.w["gu"]) + act = m3_silu_and_mul(gu, None, None, None) + out = m3_bf16_linear(act, self.w["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 + 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_hidden_states( + self, spec_metadata, layer_id: int, hidden_states, residual=None + ) -> None: + """The Eagle3 capture with this step's capture norm; 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``).""" + _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: + """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 + ) + ) + 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, ) - return times - def _select_target(self, attn_metadata: AttentionMetadata, input_ids: Optional[torch.Tensor]): - """Which target runs this step: host-side counts only, fixed per CUDA-graph capture.""" + 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 ( - attn_metadata.num_contexts == 0 - and input_ids is not None - and input_ids.ndim == 1 - and dispatch.is_decode_step(input_ids.numel()) + self.guided_decoder is not None + or self.mapping is None + or self.mapping.enable_attention_dp ): - return self._decode - return self._prefill + 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 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, - **kwargs, - ) -> torch.Tensor: - 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") - target = ( - self._select_target(attn_metadata, input_ids) - if inputs_embeds is None - else self._prefill - ) - return target.forward( - attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, **kwargs - ) + 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 not self._m3_sampled_verify_ok(spec_metadata): + return False + self._m3_target_lm_head = lm_head + return True -class PrefillTarget: - """The general case: context rows, mixed steps and wide generation steps, on the generic layer path.""" + 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.""" + return self._m3_structural_ok(attn_metadata, spec_metadata) and draft_model is not None - def __init__(self, core: MiniMaxM3Tp4ModelingV2) -> None: - self.core = core + 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 forward( - self, attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, **kwargs + def _m3_first_drafter_inputs( + self, + position_ids, + accepted_tokens, + num_tokens: int, + num_contexts: int, + spec_metadata, + draft_model: Optional[M3Eagle3Draft], ): - core = self.core - hidden_states = core.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds - residual = None - for layer in core.layers: - hidden_states, residual = layer( - position_ids=position_ids, - hidden_states=hidden_states, - attn_metadata=attn_metadata, - residual=residual, - spec_metadata=spec_metadata, - **kwargs, + """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(), ) - return hidden_states + 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) + 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_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 + ) + self._m3_fc_hidden = draft_model.fc(hs) + return aux -class DecodeTarget: - """A pure generation step of at most 48 tokens. - - 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. ``dispatch.py`` lists the row - ranges and what a covered step that cannot take its path does. - """ + 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. - def __init__(self, core: MiniMaxM3Tp4ModelingV2) -> None: - self.core = core - self._moe_push_descs: Dict[tuple, Any] = {} + _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 forward( - self, attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, **kwargs - ): - core = self.core - num_tokens = int(input_ids.numel()) - plan = plan_for(num_tokens) - hidden_states, residual = self._embed(input_ids, num_tokens) - for layer in core.layers: - if residual is None or ( - spec_metadata is not None and spec_metadata.is_layer_capture(layer.layer_idx) - ): - hidden_states, residual = layer( - position_ids=position_ids, - hidden_states=hidden_states, - attn_metadata=attn_metadata, - residual=residual, - spec_metadata=spec_metadata, - **kwargs, - ) - else: - hidden_states, residual = self._layer( - layer, plan, position_ids, hidden_states, residual, attn_metadata, kwargs - ) - return hidden_states + 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. - 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 dispatch.EMBED_FIRSTNORM.covers(num_tokens): - dispatch.uncovered( - dispatch.EMBED_FIRSTNORM.path, f"more than {EMBED_FIRSTNORM_MAX_TOKENS} tokens" - ) - return core.embed_tokens(input_ids), None - owner = core._embed_owners.get((num_tokens, input_ids.dtype)) - if owner is None or not input_ids.is_contiguous(): - dispatch.missed( - dispatch.EMBED_FIRSTNORM.path, - num_tokens, - f"no embedding owner for {input_ids.dtype} ids (built for int32 / int64 by build_targets)" - if owner is None - else "the input ids are not contiguous", - ) - return core.embed_tokens(input_ids), None - return m3_embedding_firstnorm( - input_ids, - owner.table, - owner.norm.weight, - owner.eps, - owner.raw, - owner.output, - owner.use_pdl, - ) + 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 + 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 _bucket_flags(layer: M3DecoderLayer, num_tokens: int, hidden: int): - """(bucket workspace, its 8-CTA call counters) of the boundary all-reduce for ``num_tokens``, or None.""" - workspace = m3_mc_allreduce.get_workspace(layer.mapping, 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) + 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 - @staticmethod - def _no_bucket(path: str, num_tokens: int) -> None: - if torch.cuda.is_current_stream_capturing(): - dispatch.capture_missed(path, num_tokens, "the multicast all-reduce workspace") + def _m3_target_accept( + self, + hidden_states, + spec_metadata, + batch_size: int, + draft_len: int, + draft_model: Optional[M3Eagle3Draft] = None, + ): + """Fused verification (see tail_eligible); same outputs as the strict acceptance path.""" + 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, + ) + 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 + ) + 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 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: + logits = m3_target_head_kernel.logits(head, rows) else: - dispatch.missed(path, num_tokens, "no multicast all-reduce bucket holds the step") + 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 _attn_push_target(self, layer: M3DecoderLayer, num_tokens: int, hidden: int): - """The o_proj's push tuple into the attention boundary's all-reduce, or None (after a miss). + 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 - 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(layer, num_tokens, hidden) - if found is None: - self._no_bucket(dispatch.ATTN_PUSH.path, num_tokens) + 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 - 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, layer: M3DecoderLayer, num_tokens: int, hidden: int): - """(engine push descriptor, shared down-projection push tuple) for the MoE boundary's exchange sets, or None - (after a miss). + 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 - 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 - (``M3MoE.prebuild_engines``), so it does.""" - from .ops.moe.m3_moe_engine import ar_push_desc + def _m3_force_vec( + self, batch_size: int, base_total: int, max_total: int, device + ) -> Optional[torch.Tensor]: + """Per-request forced accepted counts for a fractional forced acceptance, else None. - found = self._bucket_flags(layer, num_tokens, hidden) - if found is None: - self._no_bucket(dispatch.MOE_PUSH.path, num_tokens) + Draws exactly as _apply_force_accepted_tokens does for a decode-only batch (same fixed-seed pool, one counter + bump per step, the same (counter, slot) hash), so every TP rank and the generic path agree on the counts. + """ + frac = self.force_num_accepted_tokens - int(self.force_num_accepted_tokens) + if frac <= 0.0 or base_total >= max_total: 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(): - dispatch.capture_missed( - dispatch.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 + self._ensure_force_accept_rng_state(device) + out = torch.empty(batch_size, dtype=torch.int32, device=device) + 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()), ) - return desc, ( - set_b, - flags.data_ptr(), - workspace.rank, - workspace.world_size, - workspace.max_tokens, - 2, + + # ---------------------------------------------------------------------------------------------- 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 _boundary( - self, layer: M3DecoderLayer, use_mc: bool, x, residual, norm, x2=None, quant=False - ): - """All-reduce + (second addend) + residual + RMSNorm: the multicast kernel, else the built-in all-reduce.""" - if use_mc: - workspace = m3_mc_allreduce.get_workspace(layer.mapping, x.shape[-1]) - fused = ( - None - if workspace is None - else m3_mc_allreduce.m3_mc_allreduce_norm( - x, - residual.contiguous(), - norm.weight, - norm.variance_epsilon, - workspace, - use_pdl=get_env_enable_pdl(), - x2=x2, - quant=quant, - ) - ) - if fused is not None: - return fused - dispatch.missed( - dispatch.MC_ALLREDUCE.path, - x.shape[0], - m3_mc_allreduce.decline_reason(x, residual, workspace, x2=x2), + 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_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" ) - if x2 is not None: - x = x.add_(x2) - return layer.native_allreduce_norm(x, residual, norm) - - def _prepushed(self, layer, x, residual, norm, quant, prepushed, x2=None): - workspace = m3_mc_allreduce.get_workspace(layer.mapping, x.shape[-1]) - fused = m3_mc_allreduce.m3_mc_allreduce_norm( - x, - residual.contiguous(), - norm.weight, - norm.variance_epsilon, - workspace, - use_pdl=get_env_enable_pdl(), - x2=x2, - quant=quant, - prepushed=prepushed, - ) - if fused is None: - raise RuntimeError( - f"layer {layer.layer_idx}: rows were pushed into an all-reduce that did not run" + 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 fused + 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 _layer( + def _forward_draft_loop( self, - layer: M3DecoderLayer, - plan: DecodePlan, - position_ids, - hidden_states, - residual, + inputs, attn_metadata, - kwargs, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, ): - num_tokens = hidden_states.shape[0] - hidden = hidden_states.shape[-1] - # 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: - dispatch.uncovered( - "decode_graph", - "an eager decode step: the decode-graph GEMV dispatch and the boundary pushes are off", + """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, ) - push = plan.attn_mc_ar and decode_graph and dispatch.ATTN_PUSH.covers(num_tokens) - armed = None - if push: - armed = self._attn_push_target(layer, num_tokens, hidden) - m3_prequant.AR_PUSH.arm(armed) - hidden_states = layer.self_attn( - position_ids=position_ids, - hidden_states=hidden_states, - attn_metadata=attn_metadata, - all_reduce_params=AllReduceParams(enable_allreduce=False), - **kwargs, + 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, ) - if push: - m3_prequant.AR_PUSH.target = None - if push and m3_prequant.AR_PUSH.take_pushed(hidden_states): - hidden_states, residual = self._prepushed( - layer, hidden_states, residual, layer.post_attention_layernorm, False, AR_PUSH_MODE - ) - else: - if armed is not None: - dispatch.missed( - dispatch.ATTN_PUSH.path, - num_tokens, - "the o_proj did not push its rows into the armed boundary all-reduce", - ) - hidden_states, residual = self._boundary( - layer, plan.attn_mc_ar, hidden_states, residual, layer.post_attention_layernorm - ) - moe = layer.block_sparse_moe - if moe is None: - out = layer.mlp( - hidden_states, final_all_reduce_params=AllReduceParams(enable_allreduce=False) - ) - return self._boundary(layer, plan.attn_mc_ar, out, residual, layer.next_layer_layernorm) + 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``): - norm = layer.next_layer_layernorm - quant = plan.qkv_prequant and layer.next_is_decoder_layer - path = dispatch.moe_path(num_tokens) - moe_push = None - if path == dispatch.MOE_ENGINE.path and moe.shared_experts is not None: - # The engine itself runs inside the MoE forward; looked up (or found missing) here first so a push is - # armed only for an engine that will take it. - engine = moe.engine(num_tokens, shared=True) - if engine is not None and plan.moe_mc_ar and decode_graph: - moe_push = self._moe_push_targets(layer, num_tokens, hidden) - if moe_push is not None: - m3_prequant.MOE_PUSH.arm(*moe_push) - elif path == dispatch.MOE16_ENGINE.path and decode_graph: - # Router, routed and shared experts in one kernel; its bf16 sum goes to the boundary all-reduce. - engine = moe.moe16_engine(num_tokens) - if engine is not None: - out = m3_moe16_engine(hidden_states.contiguous(), engine) - return self._boundary(layer, plan.moe_mc_ar, out, residual, norm, quant=quant) - if ( - path in (dispatch.MOE16_ENGINE.path, dispatch.MC_FINALIZE.path) - and plan.moe_mc_ar - and moe.shared_experts is not None - ): - if self._bucket_flags(layer, num_tokens, hidden) is None: - self._no_bucket(dispatch.MC_FINALIZE.path, num_tokens) - else: - shared, fc2, scales, perm = moe( - hidden_states, - attn_metadata, - final_all_reduce_params=AllReduceParams(enable_allreduce=False), - do_finalize=False, - ) - 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.weight, - norm.variance_epsilon, - m3_mc_allreduce.get_workspace(layer.mapping, hidden), - use_pdl=get_env_enable_pdl(), - quant=quant, + * 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, ) - if fused is None: + h = draft_model.attention(x, md, reduce=False) + h, residual = draft_model.ar_norm(h, residual, "post_norm") + h = draft_model.mlp(h, reduce=False) + h, to_save = draft_model.ar_norm(h, residual, "final_norm") + 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( - "the MC finalize all-reduce rejected the unfinalized MoE outputs" + "M3 fused draft step unavailable after _m3_fused_loop_ok() accepted the step" ) - return fused - parts = moe( - hidden_states, - attn_metadata, - final_all_reduce_params=AllReduceParams(enable_allreduce=False), - return_parts=True, + 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" ) - shared, routed = parts if isinstance(parts, tuple) else (parts, None) - if moe_push is not None: - m3_prequant.MOE_PUSH.disarm() - if routed is None: - # The engine folded the shared expert in and pushed routed + shared into set A. - folded_pushed = m3_prequant.MOE_PUSH.engine_pushed == shared.data_ptr() - m3_prequant.MOE_PUSH.engine_pushed = m3_prequant.MOE_PUSH.shared_pushed = None - if folded_pushed: - return self._prepushed(layer, shared, residual, norm, quant, 2) - else: - routed_pushed, shared_pushed = m3_prequant.MOE_PUSH.take_pushed(shared, routed) - if routed_pushed and shared_pushed: - return self._prepushed(layer, shared, residual, norm, quant, 3, x2=routed) - if routed_pushed or shared_pushed: - raise RuntimeError("only one MoE part pushed into the boundary all-reduce") - dispatch.missed( - dispatch.MOE_PUSH.path, - num_tokens, - "the MoE did not push its output into the armed boundary all-reduce", - ) - return self._boundary(layer, plan.moe_mc_ar, shared, residual, norm, x2=routed, quant=quant) + 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, all measured on GB300 tp4 AgentX decode. + + * ``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. + * 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. + """ + return RuntimeHooks( + pack_step_inputs=True, + skip_unchanged_h2d=True, + loop_on_execution_stream=True, + merge_tuned_tactics_across_ranks=True, + skip_greedy_rng_upload=True, + spec_step_fold=spec_fold, + msa=MsaRuntimePolicy( + short_sequence_identity=True, + kv_update=_msa_kv_update_kernel.msa_kv_update, + skip_unchanged_block_offsets=True, + reuse_pool_views=True, + ), + ) -class ModelingV2MiniMaxM3Nvfp4Sm103Tp4(M3SpecDecForwardMixin, MiniMaxM3ForCausalLM): +class ModelingV2MiniMaxM3Nvfp4Sm103Tp4( + DecoderModelForCausalLM[MiniMaxM3Tp4ModelingV2, PretrainedConfig] +): """MiniMax-M3 NVFP4 / sm_103 / tp4, text only, with one-model Eagle3 or no speculation. - A subclass of the built-in text model shell for its loader, its KV-cache and - transceiver preferences and its one-engine speculative plumbing; the text - model it wraps, the forward and the speculative worker are this target's. + 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: @@ -570,64 +4532,132 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: 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 (dispatch.py). - dispatch.set_require(resolved_mode(type(self)) is ModelingV2Mode.REQUIRE) - if is_minimax_m3_vl_config(model_config.pretrained_config): - model_config = get_text_model_config(model_config) + # 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 - # The one-engine shell's constructor, past the built-in one that builds the built-in text model. - SpecDecOneEngineForCausalLM.__init__( - self, MiniMaxM3Tp4ModelingV2(model_config), model_config + super().__init__( + MiniMaxM3Tp4ModelingV2(model_config), + config=model_config, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, ) - self.runtime_hooks, self._spec_fold = make_runtime_hooks(self) - if self.spec_worker is not None: - self.spec_worker = M3Eagle3Worker.adopt(self.spec_worker) - worker = self.spec_worker - for layer in self.model.layers: - layer.capture = worker.capture_hidden_states + 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) + spec_fold = M3SpecStepFold() + self.runtime_hooks = _runtime_hooks(spec_fold) def load_weights(self, weights: Dict, *args, **kwargs) -> None: - _weights.load(self, weights, *args, **kwargs) + _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 setup_aliases(self) -> None: - """Chain each decoder layer's next_layer_layernorm for the fused boundaries.""" - super().setup_aliases() - layers = self.model.layers - for idx, layer in enumerate(layers): - layer.next_is_decoder_layer = idx < len(layers) - 1 + 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 fused verify's coded target head and the fused draft loop's head copy.""" super().cache_derived_state() - self.model.build_targets() + self.model.post_load_weights() if self.spec_worker is not None: - # The fused verify's coded target head, from the loaded weights, before the executor's warmup. - start = time.perf_counter() - if self.spec_worker.prebuild_target_head(self.lm_head): - torch.cuda.synchronize() - logger.info( - "[M3 prebuild] coded target head built before warmup: %.1f s" - % (time.perf_counter() - start) - ) - # The fused draft loop's head copy, from the loaded draft weights, before the executor's warmup. - start = time.perf_counter() - if self.spec_worker.prebuild_draft_head(self.draft_model): - torch.cuda.synchronize() - logger.info( - "[M3 prebuild] fused draft head built before warmup: %.1f s" - % (time.perf_counter() - start) - ) + self.spec_worker.prebuild(self.lm_head, self.draft_model) - def forward(self, attn_metadata: AttentionMetadata, *args, **kwargs): - """The fused-verify one-engine forward (``M3SpecDecForwardMixin``), text only.""" - multimodal_params: List[Any] = kwargs.pop("multimodal_params", None) or [] - if any( - param is not None and getattr(param, "multimodal_data", None) - for param in multimodal_params - ): + 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" ) - return super().forward(attn_metadata, *args, **kwargs) + 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, + **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) + return 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, + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/activation/__init__.py similarity index 60% rename from tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/activation/__init__.py index d6a67225eb01..b8c7426e3933 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/activation/__init__.py @@ -1,3 +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.""" +"""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_103__tp4/ops/activation/m3_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/activation/m3_silu_and_mul.py new file mode 100644 index 000000000000..64cb00558760 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp4/ops/attention/m3_draft_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_draft_attention.py new file mode 100644 index 000000000000..6347348d1681 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..bcd928fccf5b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp4/ops/attention/m3_msa_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_msa_attention.py new file mode 100644 index 000000000000..dac983d8982b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp4/ops/attention/m3_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..6493c4a45025 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp4/ops/attention/m3_qkv_insert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_qkv_insert.py new file mode 100644 index 000000000000..91927ab0bc39 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp4/ops/attention/m3_select_blocks.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_select_blocks.py new file mode 100644 index 000000000000..6bf1764bfc0d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp4/ops/comm/m3_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_allreduce.py new file mode 100644 index 000000000000..1a275c4ad863 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp4/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/ops/embedding/m3_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding/m3_embedding.py new file mode 100644 index 000000000000..f979fdecf376 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp4/ops/gemm/m3_bf16_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_bf16_linear.py new file mode 100644 index 000000000000..d74fc709dbdf --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp4/ops/gemm/m3_fp32_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_fp32_linear.py new file mode 100644 index 000000000000..8d7374810675 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp4/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv.py index 6a454a7d2af5..5c1b2ada2d6d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv.py @@ -9,8 +9,8 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.mx_gemv import op as _gemv +from ...kernels.mx_gemv import op as _gemv from ..moe.m3_moe_engine import mxfp8_dequant __all__ = ["PROJECTIONS", "gemv_is_valid", "m3_mx_gemv", "served_cells"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv_bf16.py index 05304eec0953..0951d082793a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv_bf16.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv_bf16.py @@ -9,8 +9,8 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.mx_gemv import op as _gemv +from ...kernels.mx_gemv import op as _gemv from ..moe.m3_moe_engine import mxfp8_quant_dequant from .m3_mx_gemv import ( DENSE_DOWN, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_gemm.py new file mode 100644 index 000000000000..b3d54e14ef2e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + # 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_103__tp4/ops/gemm/m3_mxfp8_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mxfp8_linear.py new file mode 100644 index 000000000000..799cda007837 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp4/ops/moe/m3_moe16_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe16_engine.py index a643aa0de12c..0e1595ed8ccf 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe16_engine.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe16_engine.py @@ -11,9 +11,9 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.moe.op import M3MoeLayer -from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.moe.op16 import M3Moe16State +from ...kernels.m3_moe.op import M3MoeLayer +from ...kernels.m3_moe.op16 import M3Moe16State from .m3_moe_engine import moe_reference __all__ = ["M3Moe16State", "m3_moe16_engine"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe_engine.py index d8167bb53c0b..219537a853f4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe_engine.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe_engine.py @@ -10,7 +10,12 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.moe.op import ( +from tensorrt_llm.quantization.utils.fp4_utils import ( + get_reorder_rows_for_gated_act_gemm_row_indices, + get_shuffle_matrix_a_row_indices, +) + +from ...kernels.m3_moe.op import ( HIDDEN, INTER, SHARED_INTER, @@ -20,10 +25,6 @@ ar_push_desc, is_supported, ) -from tensorrt_llm.quantization.utils.fp4_utils import ( - get_reorder_rows_for_gated_act_gemm_row_indices, - get_shuffle_matrix_a_row_indices, -) __all__ = [ "M3MoeLayer", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/m3_embedding_firstnorm.py index ff7fb16e2266..535f928018f5 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/m3_embedding_firstnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/m3_embedding_firstnorm.py @@ -9,10 +9,8 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.embedding_firstnorm.op import ( - DecodeEmbedding, - m3_embed_norm, -) + +from ...kernels.embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm __all__ = ["DecodeEmbedding", "m3_embedding_firstnorm"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_force_vec.py new file mode 100644 index 000000000000..1905516a8403 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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 ...kernels.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_103}) + + 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_103__tp4/runtime.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/runtime.py deleted file mode 100644 index c40314bf2302..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/runtime.py +++ /dev/null @@ -1,50 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The step-loop hooks this target opts into (``RuntimeHooks``), all measured on GB300 tp4 AgentX decode. - -* ``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. -* 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. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Optional, Tuple - -from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy, RuntimeHooks - -from .kernels import msa_kv_update -from .spec_fold import M3SpecStepFold - -if TYPE_CHECKING: - from .modeling import ModelingV2MiniMaxM3Nvfp4Sm103Tp4 - - -def make_runtime_hooks( - model: "ModelingV2MiniMaxM3Nvfp4Sm103Tp4", -) -> Tuple[RuntimeHooks, Optional[M3SpecStepFold]]: - """This target's hooks, and the speculative step fold they carry (None without speculation).""" - spec_fold = M3SpecStepFold() if model.spec_config is not None else None - hooks = RuntimeHooks( - pack_step_inputs=True, - skip_unchanged_h2d=True, - loop_on_execution_stream=True, - merge_tuned_tactics_across_ranks=True, - skip_greedy_rng_upload=True, - spec_step_fold=spec_fold, - msa=MsaRuntimePolicy( - short_sequence_identity=True, - kv_update=msa_kv_update.msa_kv_update, - skip_unchanged_block_offsets=True, - reuse_pool_views=True, - ), - ) - return hooks, spec_fold diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec_fold.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec_fold.py deleted file mode 100644 index 593a340a2782..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec_fold.py +++ /dev/null @@ -1,112 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""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 never stores the sampler's outputs itself, so ``begin_step`` and -``sampler_rows_stored`` always report that nothing was staged. - -* ``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``). -""" - -from __future__ import annotations - -from typing import Optional - -import torch - -from .kernels import m3_spec_tail - - -class M3SpecStepFold: - """``SpecStepFold`` implementation of the MiniMax-M3 tp4 target (see the module docstring). - - One instance per target model; the target passes it to ``RuntimeHooks.spec_step_fold``. - """ - - def __init__(self) -> None: - self.store: Optional[object] = None - self.step_inputs: Optional[ - tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] - ] = 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) - - def bind_sampler_store(self, store: object) -> None: - self.store = store - - def begin_step(self, seq_slots: Optional[list[int]], num_rows: int) -> bool: - return False - - 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: - return False - - 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, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec_worker.py deleted file mode 100644 index 773fdb9aa798..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec_worker.py +++ /dev/null @@ -1,1166 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Eagle3 one-model speculative decoding of the MiniMax-M3 tp4 target, at stock numerics. - -``M3Eagle3Worker`` is the in-tree ``Eagle3OneModelWorker`` with a fused verify and a fused draft loop: - -* Verify: on decode-only batches whose verify rows the fused kernels cover, the target's BF16 lm_head computes this - rank's vocabulary shard of the logits (the in-tree LM head's own local GEMM; at most 8 verify rows stream a - lossless 12-bit coding of the same BF16 weight, ``dispatch.TARGET_HEAD``) and one kernel replaces the logits - all-gather, the argmax (or the temperature + top-p draw) and the strict acceptance. -* Draft: the in-tree Eagle3 drafter's modules and weights at their precision (BF16 weights and activations, the - full-vocabulary BF16 head, greedy drafts). The first pass may take the per-layer fc_norm RMSNorms in one launch (or - folded into the hidden-state capture) and run its fc beside the target head GEMM (at most 8 verify rows: forked - right after the last captured target layer, ``dispatch.FC_EARLY_MAX_ROWS``). On a decode-only step of at most - 8 requests (``dispatch.DRAFT_STEP``) the draft loop is fused (``_m3_fused_draft_loop``): each pass's draft input in - one launch, the drafter layer's in-tree 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. - -The target's outer CausalLM drives the worker through ``M3SpecDecForwardMixin.forward`` (or the same calls made from -its own forward), and its layer loop captures the Eagle3 hidden states through ``capture_hidden_states``. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any, Optional - -import torch -import triton -import triton.language as tl -from torch import nn - -from tensorrt_llm._torch.distributed import AllReduce, AllReduceParams -from tensorrt_llm._torch.models.modeling_speculative import _slice_spec_position_ids -from tensorrt_llm._torch.modules.gated_mlp import GatedMLP -from tensorrt_llm._torch.modules.linear import TensorParallelMode -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager -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, -) -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping - -from . import dispatch -from .kernels import m3_draft_step, m3_sampled_accept, m3_shard_accept, m3_spec_tail, m3_target_head - -if TYPE_CHECKING: - from tensorrt_llm._torch.model_config import ModelConfig - from tensorrt_llm.llmapi.llm_args import EagleDecodingConfig - -# Id of the target's final hidden state in ``layers_to_capture`` (SpecDecOneEngineForCausalLM.layer_idx). -FINAL_HIDDEN_LAYER_ID = -1 - - -@triton.jit -def _m3_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) - - -def draft_ar_norm( - all_reduce: AllReduce, x: torch.Tensor, residual: torch.Tensor, norm: nn.Module -) -> tuple[torch.Tensor, torch.Tensor]: - """One draft-layer boundary in stock Eagle3's structure: the projection's plain all-reduce (BF16 out), then the - norm module with the residual (FlashInfer's fused_add_rmsnorm: FP32 add, norm of the FP32 sum). Returns the - normalized rows and the residual sum.""" - return norm(all_reduce(x), residual) - - -def capture_hidden_states( - spec_metadata: Any, - layer_id: int, - hidden_states: torch.Tensor, - residual: Optional[torch.Tensor], - capture_norm: Optional[tuple[torch.Tensor, float]], -) -> 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) (``M3Eagle3Worker.capture_norm``), 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 " - f"buffer capacity: num_tokens={num_tokens}, " - f"capacity={spec_metadata.hidden_states.shape[0]}" - ) - if num_tokens > hidden_states.shape[0]: - raise RuntimeError( - "EAGLE3 hidden-state capture token count exceeds " - f"available hidden states: num_tokens={num_tokens}, " - f"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 - - -class M3Eagle3Worker(Eagle3OneModelWorker): - """``Eagle3OneModelWorker`` with the MiniMax-M3 fused verify (see the module docstring). - - Methods whose docstring says "Mirrors" are copies of the named in-tree method with the fused-verify branches - inserted; compare them with that method when it changes. - """ - - def __init__( - self, - spec_config: "EagleDecodingConfig", - mapping: Optional[Mapping] = None, - model_config: Optional["ModelConfig"] = None, - use_separate_draft_kv_cache: bool = False, - ): - super().__init__( - spec_config, - mapping=mapping, - model_config=model_config, - use_separate_draft_kv_cache=use_separate_draft_kv_cache, - ) - self._init_fused_verify() - - @classmethod - def adopt(cls, worker: Eagle3OneModelWorker) -> "M3Eagle3Worker": - """Turn the in-tree worker the one-engine shell built (``get_spec_worker``) into this worker, in place.""" - if type(worker) is not Eagle3OneModelWorker: - raise TypeError(f"expected an Eagle3OneModelWorker, got {type(worker).__name__}") - worker.__class__ = cls - worker._init_fused_verify() - return worker - - def _init_fused_verify(self) -> None: - # 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_norm[0].weight, stacked fc_norm weights [layers, H]), built eagerly on the first fused first pass. - self._m3_fc_norm_w: Optional[tuple[torch.Tensor, torch.Tensor]] = None - # Armed per step by arm_capture_norm: (stacked fc_norm weights, eps) or None. - self._m3_capture_norm: Optional[tuple[torch.Tensor, float]] = 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 - # (capture_hidden_states), and the fork that did it, joined by _m3_target_accept. - self._m3_fc_early_drafter: Optional[nn.Module] = 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 at <= 8 rows, built before the - # executor's warmup (prebuild_target_head). - self._m3_target_head: Optional[m3_target_head.TargetHead] = None - self._m3_force_vec_warm = False - # The fragment-ordered copy of the drafter's BF16 head shard the fused draft loop streams, built before the - # executor's warmup (prebuild_draft_head). - self._m3_draft_head: Optional[m3_draft_step.DraftHead] = 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 - - @property - def capture_norm(self) -> Optional[tuple[torch.Tensor, float]]: - """What ``capture_hidden_states`` needs for this step's captures (set by ``arm_capture_norm``).""" - return self._m3_capture_norm - - def capture_hidden_states( - self, - spec_metadata: Any, - layer_id: int, - hidden_states: torch.Tensor, - residual: Optional[torch.Tensor] = None, - ) -> None: - """``capture_hidden_states`` with this step's capture norm; 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``).""" - capture_hidden_states( - 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 _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 - ) - ) - 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 because - # the accepted token count affects which Mamba states are valid. The - # isinstance check below naturally no-ops on non-Mamba kv_cache_managers, - # so this is safe to run unconditionally regardless of spec mode (Eagle3 - # over a Mamba-style draft is plausible, even if no such draft exists today). - 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: torch.LongTensor, - position_ids: torch.LongTensor, - hidden_states: torch.Tensor, - accepted_tokens: torch.Tensor, - attn_metadata: Any, - spec_metadata: Any, - draft_model: nn.Module, - ): - """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: Any, spec_metadata: Any) -> 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) > dispatch.DECODE_MAX_TOKENS: - return False - return not spec_metadata.wants_advanced_draft_sampling - - def tail_eligible(self, attn_metadata: Any, spec_metadata: Any, 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 ``dispatch.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 self._m3_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 dispatch.SHARD_ACCEPT.covers(key): - dispatch.uncovered( - dispatch.SHARD_ACCEPT.path, - f"draft length {draft_len} 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: Any, spec_metadata: Any, draft_model: Optional[nn.Module] - ) -> 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.""" - return self._m3_structural_ok(attn_metadata, spec_metadata) and self._m3_fc_norm_ok( - draft_model - ) - - def _m3_seed_first_pass( - self, num_accepted_tokens: torch.Tensor, spec_metadata: Any, 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: torch.Tensor, - accepted_tokens: torch.Tensor, - num_tokens: int, - num_contexts: int, - spec_metadata: Any, - draft_model: nn.Module, - ) -> Optional[dict[str, Any]]: - """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 in-tree 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 not self._m3_fc_norm_ok(draft_model) - ): - return None - dm = draft_model.model - layers = len(spec_metadata.layers_to_capture) - if dm.hidden_size_in != spec_metadata.hidden_size or layers != len(dm.fc_norm): - return None - cache = self._m3_fc_norm_w - if cache is None or cache[0] is not dm.fc_norm[0].weight: - if torch.cuda.is_current_stream_capturing(): - dispatch.capture_missed("spec_tail", None, "the stacked Eagle3 fc_norm weights") - return None - cache = ( - dm.fc_norm[0].weight, - torch.stack([n.weight.detach() for n in dm.fc_norm]).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], dm.fc_norm[0].variance_epsilon) - hidden = dm.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} - - @staticmethod - def _m3_fc_norm_ok(draft_model: Optional[nn.Module]) -> bool: - """An Eagle3 drafter with an fc over several captured layers, each normalized by its own fc_norm (one - shared epsilon): what the fused first drafter inputs compute.""" - dm = getattr(draft_model, "model", None) - fc_norm = getattr(dm, "fc_norm", None) - return ( - fc_norm is not None - and getattr(dm, "fc", None) is not None - and len({n.variance_epsilon for n in fc_norm}) == 1 - ) - - def _m3_fc_fork( - self, - spec_metadata: Any, - num_tokens: int, - draft_model: Optional[nn.Module], - 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 or not self._m3_fc_norm_ok(draft_model): - return None - dm = draft_model.model - layers = len(spec_metadata.layers_to_capture) - if ( - dm.hidden_size_in != spec_metadata.hidden_size - or layers != len(dm.fc_norm) - 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], dm.fc_norm[0].variance_epsilon, use_pdl=False - ) - self._m3_fc_hidden = dm.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 or - torch.compile, 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: Any, - spec_metadata: Any, - lm_head: nn.Module, - draft_model: Optional[nn.Module], - ) -> 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 - ``dispatch.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 - # 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) - <= dispatch.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 self._m3_fc_norm_ok(draft_model) - and fc_norm_w[0] is draft_model.model.fc_norm[0].weight - and len(spec_metadata.layers_to_capture) == fc_norm_w[1].shape[0] - ): - self._m3_capture_norm = (fc_norm_w[1], draft_model.model.fc_norm[0].variance_epsilon) - - @staticmethod - def _m3_sampled_verify_ok(spec_metadata: Any) -> 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: torch.Tensor, - spec_metadata: Any, - batch_size: int, - draft_len: int, - draft_model: Optional[nn.Module] = None, - ): - """Fused verification (see tail_eligible); same outputs as the strict acceptance path.""" - 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_kernel[(1,)]( - 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=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 - ) - 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 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 in-tree LM 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: - logits = m3_target_head.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.shard_accept( - logits, - vocab_offset, - self.mapping, - draft_tokens, - batch_idx, - force_total, - self._m3_next_new, - force_vec=force_vec, - ) - else: - n_rows = rows.shape[0] - seeds, offsets = self._rng_state_per_token(spec_metadata, n_rows) - result = m3_sampled_accept.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, - ) - 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: - raise RuntimeError( - "M3 fused verify unavailable after tail_eligible() accepted the step" - ) - 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 <= 8 rows (rebuilt from - the current weight values on every call). Called by the target's ``cache_derived_state`` once the weights are - real, before the executor's warmup, 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 in-tree GEMM runs.""" - self._m3_target_head = None - if not self._m3_sharded_head(lm_head) or not m3_target_head.head_supported(lm_head): - return False - head = m3_target_head.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 in-tree BF16 head GEMM" - ) - return head.coded - - def _m3_coded_head(self, lm_head: nn.Module, rows: int) -> Optional[m3_target_head.TargetHead]: - """The coded target head for a fused verify of ``rows`` rows, or None for the in-tree GEMM: above - ``dispatch.TARGET_HEAD_MAX_ROWS`` rows (uncovered), on a shard whose coding failed its round trip - (uncovered: the in-tree GEMM is the exact path), and when the copy was not built for this weight (a miss).""" - path = dispatch.TARGET_HEAD.path - if not dispatch.TARGET_HEAD.covers(rows): - dispatch.uncovered(path, f"{rows} verify rows (the in-tree head GEMM runs)") - return None - head = self._m3_target_head - if head is None or head.weight is not lm_head.weight: - dispatch.missed( - path, - rows, - "the coded target head is not built for the head's current weight (prebuild_target_head)", - ) - return None - if not head.coded: - dispatch.uncovered( - path, "the coded shard failed its bitwise round trip (the in-tree head GEMM runs)" - ) - return None - return head - - def _m3_force_vec( - self, batch_size: int, base_total: int, max_total: int, device: torch.device - ) -> Optional[torch.Tensor]: - """Per-request forced accepted counts for a fractional forced acceptance, else None. - - Draws exactly as _apply_force_accepted_tokens does for a decode-only batch (same fixed-seed pool, one counter - bump per step, the same (counter, slot) hash), so every TP rank and the generic path agree on the counts. - """ - 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) - # One launch: bump the shared counter and draw, with _apply_force_accepted_tokens' exact arithmetic. - _m3_force_vec_kernel[(1,)]( - 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=max(16, 1 << (batch_size - 1).bit_length()), - ) - return out - - @staticmethod - def _m3_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 - ) - - # ---------------------------------------------------------------------------------------------- fused draft loop - - def prebuild_draft_head(self, draft_model: Optional[nn.Module]) -> bool: - """Build the fragment-ordered copy of the drafter's BF16 head shard that the fused draft loop streams (rebuilt - from the current weight values on every call). Called by the target's ``cache_derived_state`` once the weights - are real, before the executor's warmup, 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.prepare_head(draft_model.lm_head, self.mapping.tp_rank) - return True - - def _m3_drafter_supported(self, draft_model: nn.Module) -> bool: - """The drafter the fused draft loop runs: one plain (non-MLA) Eagle3 layer with a gated MLP over the - target-vocabulary embedding (whole table on every rank), a draft vocabulary equal to the target's (no d2t), - all-reduced TP projections, and a bias-free vocab-sharded BF16 head without padding.""" - dm = getattr(draft_model, "model", None) - if ( - dm is None - or getattr(dm, "num_layers", 0) != 1 - or dm._use_mla - or dm._eh_proj_before_attn - ): - return False - layer = dm.midlayer - if layer._next_layer_regular or not isinstance(layer.mlp, GatedMLP): - return False - embed = dm.embed_tokens - if embed is None or getattr(embed, "tp_size", 1) != 1: - return False - if self.mapping is None or self.mapping.tp_size <= 1 or self.mapping.enable_attention_dp: - return False - if getattr(layer.self_attn.o_proj, "all_reduce", None) is None: - return False - if getattr(layer.mlp.down_proj, "all_reduce", None) is None: - return False - if self._d2t is not None or getattr(dm, "d2t", None) is not None: - return False - lm_head = draft_model.lm_head - return self._m3_sharded_head(lm_head) and m3_draft_step.head_supported(lm_head) - - def _m3_fused_loop_ok( - self, - attn_metadata: Any, - spec_metadata: Any, - batch_size: int, - draft_model: Optional[nn.Module], - ) -> 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) - ``dispatch.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 = dispatch.DRAFT_STEP.path - if not self._m3_drafter_supported(draft_model) or not hasattr( - attn_metadata, "kv_lens_cuda" - ): - dispatch.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 dispatch.DRAFT_STEP.covers((draft_len, batch_size)): - dispatch.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: - dispatch.missed( - path, - batch_size, - "the drafter head's fragment-ordered copy is not built for its current weight (prebuild_draft_head)", - ) - return False - return m3_draft_step.workspace(self.mapping, head.wf.device) is not None - - 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: dict[str, Any], - attn_metadata: Any, - spec_metadata: Any, - draft_model: nn.Module, - draft_kv_cache_manager: Any, - batch_size: int, - num_accepted_tokens: torch.Tensor, - ) -> torch.Tensor: - """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_layernorm 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 in-tree attention and MLP, each followed by the in-tree all-reduce with the residual add - and the next RMSNorm fused into it (post_attention_layernorm, 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 - dm = draft_model.model - layer = dm.midlayer - attn = layer.self_attn - mlp = layer.mlp - attn_all_reduce = attn.o_proj.all_reduce - mlp_all_reduce = mlp.down_proj.all_reduce - table = dm.embed_tokens.weight - 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:] - - no_all_reduce = AllReduceParams(enable_allreduce=False) - 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, - layer.input_layernorm.weight, - layer.input_layernorm.variance_epsilon, - layer.hidden_norm.weight, - layer.hidden_norm.variance_epsilon, - ) - h = attn( - position_ids=pos, - hidden_states=x, - attn_metadata=md, - all_reduce_params=no_all_reduce, - ) - h, residual = draft_ar_norm( - attn_all_reduce, h, residual, layer.post_attention_layernorm - ) - h = mlp(h, final_all_reduce_params=no_all_reduce) - h, to_save = draft_ar_norm(mlp_all_reduce, h, residual, dm.norm) - if dm._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.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 - - -class M3SpecDecForwardMixin: - """The target CausalLM's step forward with the fused-verify hooks of the speculative wrapper. - - Mirrors ``SpecDecOneEngineForCausalLM.forward`` (tensorrt_llm/_torch/models/modeling_speculative.py) with the - fused-verify branches: before the target forward the worker arms the capture norm, and a step the worker verifies - from the logits shard (``M3Eagle3Worker.tail_eligible``) skips the gathered target logits. - - The class it is mixed into (listed before ``DecoderModelForCausalLM``) provides ``model`` (called like the in-tree - decoder model, capturing through ``M3Eagle3Worker.capture_hidden_states``), ``lm_head``, ``logits_processor``, - ``model_config``, ``draft_model`` and ``spec_worker`` (an ``M3Eagle3Worker``, or None without speculation). - """ - - def forward( - self, - attn_metadata: Any, - 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[Any] = None, - resource_manager: Optional[Any] = None, - **kwargs, - ) -> Any: - spec_worker: Optional[M3Eagle3Worker] = self.spec_worker - if spec_worker is not None and spec_metadata is not None: - spec_worker.arm_capture_norm( - attn_metadata, spec_metadata, self.lm_head, self.draft_model - ) - hidden_states = self.model( - input_ids=input_ids, - attn_metadata=attn_metadata, - position_ids=position_ids, - inputs_embeds=inputs_embeds, - spec_metadata=spec_metadata, - **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): - if spec_worker is not None: - spec_worker.capture_hidden_states( - spec_metadata, FINAL_HIDDEN_LAYER_ID, hidden_states, None - ) - else: - capture_hidden_states( - spec_metadata, FINAL_HIDDEN_LAYER_ID, hidden_states, None, None - ) - if attn_metadata.padded_num_tokens is not None: - hidden_states = hidden_states[: attn_metadata.num_tokens] - - if spec_worker is not None: - # The fused verify computes this rank's logits shard itself and needs no gathered logits. - if spec_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 - ) - return spec_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, - ) - - return self.logits_processor.forward( - hidden_states, - self.lm_head, - attn_metadata, - return_context_logits, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/weights.py index 03d9b08caa6b..da27f2fba67f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/weights.py @@ -1,62 +1,865 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Weight loading: MiniMax-M3 NVFP4 / sm_103 / tp4. - -The checkpoint is the vision-language wrapper's. ``language_model.*`` holds the -language model; ``vision_tower.*``, ``multi_modal_projector.*`` and -``patch_merge_mlp.*`` hold the vision tower and its projector. This target is -text only, so the keys split three ways: - -* ``language_model.*`` goes to the built-in MiniMax-M3 text loader with the - prefix stripped (its multi-token-prediction head, ``model.mtp.*``, is not - loaded, as in the built-in model). That loader packs the five-way - Q/K/V/index-Q/index-K projection, folds the Gemma boundary norms and shards - every tensor for tp4; the modules it fills are this target's subclasses of - the built-in ones, under the built-in names. -* The vision tower and the projector are a predicted non-load: listed here, - never read. -* Any other key fails the load, naming it, rather than being dropped. +"""Weight tables and loaders: MiniMax-M3 NVFP4 / sm_103 / 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 typing import Any, Dict, Optional +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.models.checkpoints.base_weight_mapper import BaseWeightMapper -from tensorrt_llm._torch.models.modeling_minimaxm3 import ( - MiniMaxM3ForCausalLM, - _strip_language_model_prefix, +from tensorrt_llm._torch.moe.fused_moe.quantization import ( + trtllmgen_maybe_get_cached_w2_permute_indices, + trtllmgen_maybe_get_cached_w3_w1_permute_indices, ) -_LANG_PREFIX = "language_model." +#: ``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 + -# The vision tower's and the projector's key families: in the checkpoint, not loaded by this text-only target. -PREDICTED_NON_LOAD = ("vision_tower.", "multi_modal_projector.", "patch_merge_mlp.") +def _f32(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return t.to(torch.float32) -def load( - model, - weights: Dict[str, Any], - weight_mapper: Optional[BaseWeightMapper] = None, - params_map: Optional[Dict[str, str]] = None, - allow_partial_loading: bool = False, -) -> None: - """Load the language model's weights into ``model``, the target shell, and check every other key is predicted.""" - unknown = sorted( - k - for k in weights.keys() - if not k.startswith(_LANG_PREFIX) and not k.startswith(PREDICTED_NON_LOAD) +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 ) - assert not unknown, ( - f"{len(unknown)} checkpoint key(s) are neither language-model weights nor a predicted non-load, " - f"first {unknown[:5]}" + 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) ) - text_weights, _ = _strip_language_model_prefix( - {k: v for k, v in weights.items() if k.startswith(_LANG_PREFIX)} + + +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 ) - assert text_weights, f"the checkpoint has no {_LANG_PREFIX}* keys" - MiniMaxM3ForCausalLM.load_weights( - model, - text_weights, - weight_mapper=weight_mapper, - params_map=params_map, - allow_partial_loading=allow_partial_loading, + 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_103 / 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_103__tp8/dispatch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/dispatch.py deleted file mode 100644 index 26df46630059..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/dispatch.py +++ /dev/null @@ -1,527 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The decode-step fast paths of this target: the rows each one takes, the catalog entry that certifies them, and -what happens when one cannot run. - -One target serves every concurrency with the same code; which fast path a step takes is a pure function of its -shape, written here once (``plan_for``, ``moe_path``, ``router_path``, ``attention_path``, ``linear_path``) and -called by the forward. Every row range is a named constant below, and ``FAST_PATHS`` pairs each with the catalog -entry whose ``CELLS`` certify it: ``test_modeling_v2_m3_dispatch.py`` checks that every row count (or shape key) a -range covers is one of the entry's cells. - -Two kinds of not taking a fast path, kept apart: - -* **uncovered**: the step is outside every covered range -- prefill, mixed steps, generation wider than the ranges, - an eager decode step (the GEMV dispatch is a decode-graph one), a long-context request, an Eagle3 capture layer - that runs the generic layer by design. ``uncovered`` logs the reason once per (path, reason); the caller runs the - generic path in every mode. -* **missed**: the step is inside a covered range but a precondition is missing -- a workspace, a push descriptor or - a compiled engine not built, or a kernel declining operands it is certified for. Under ``modeling_v2='require'`` - ``missed`` raises ``FastPathUnavailable`` naming the path, the rows and the precondition; otherwise it logs once - per (path, reason) and the caller falls back. - -Compiled kernels (the MoE engines, the decode GEMVs with their push variants, the index scorer) are built by -``build_targets`` 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, kernel modules) 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``) rather -than a silent capture of the fallback. -""" - -from __future__ import annotations - -import threading -from dataclasses import dataclass -from types import MappingProxyType -from typing import Any, Callable, Dict, FrozenSet, Hashable, Iterable, Mapping, Optional, Tuple - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper -from tensorrt_llm.logger import logger - -from .ops.attention.m3_dense_decode import m3_dense_decode -from .ops.attention.m3_fused_score_select import ( - MAX_REQUESTS as FUSED_SCORE_SELECT_KERNEL_MAX_REQUESTS, -) -from .ops.attention.m3_fused_score_select import m3_fused_score_select -from .ops.attention.m3_index_score import m3_index_score -from .ops.attention.m3_producer import m3_producer -from .ops.attention.m3_select_long import MAX_BLOCKS as SELECT_LONG_MAX_BLOCKS -from .ops.attention.m3_select_long import m3_select_long -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_mc_allreduce_norm import m3_mc_allreduce_norm -from .ops.comm.m3_mc_finalize_allreduce_norm import m3_mc_finalize_allreduce_norm -from .ops.gemm.m3_draft_lowm import TACTICS as DRAFT_LOWM_TACTICS -from .ops.gemm.m3_draft_lowm import m3_draft_lowm -from .ops.gemm.m3_mx_gemv import DENSE_DOWN, DENSE_GATE_UP, O_PROJ, QKV, QKV_SPARSE, m3_mx_gemv -from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 -from .ops.moe.m3_moe_engine import m3_moe_engine -from .ops.moe.m3_router_gemv import m3_router_gemv -from .ops.moe.m3_router_gemv_fp4 import m3_router_gemv_fp4 -from .ops.moe.m3_router_tc import m3_router_tc -from .ops.norm.m3_embedding_firstnorm import m3_embedding_firstnorm -from .ops.spec.m3_draft_input import m3_draft_input -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_shard_accept import m3_shard_accept - -# -------------------------------------------------------------------------------------------------------------------- -# Row ranges -# -------------------------------------------------------------------------------------------------------------------- - -# The largest generation step the decode target takes: 8 requests verifying 1 + 3 drafts and their padding rows. -DECODE_MAX_TOKENS = 48 - - -@dataclass(frozen=True) -class DecodePlan: - """The layer-boundary collectives of one decode token bucket (GB300 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 (kernels/m3_mc_allreduce.BUCKETS). -DECODE_PLANS = ( - DecodePlan("tiny", 8, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), - DecodePlan("small", 16, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), - DecodePlan("verify", DECODE_MAX_TOKENS, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), -) - -# 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. Two requests are neutral (the grid is unchanged) and keep the step's bound. -DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({1: 16384}) -# The same cap on the Eagle3 drafter layer's trtllm-gen decode (8 query heads over one KV head per rank): GB300, one -# request, 63.7K cached tokens: 22.5 -> 11.7 us at one query row, 24.3 -> 12.5 us at four; 188K tokens: 28.2 -> 20.2 us. -# Two requests are flat and keep the step's bound. -DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({1: 16384}) - -# The sparse layers' NVFP4 QKV + indexer producer (ops/attention/m3_producer, every step its geometry fits) triggers its -# dependent kernel on entry, so the kernel after it launches and stages its pre-wait loads while the producer runs. -PRODUCER_EARLY_TRIGGER = True - -# The fused embedding + first RMSNorm kernel: one owner per row count and index dtype. -EMBED_FIRSTNORM_MAX_TOKENS = 8 -# 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 = 16 -AR_PUSH_MODE = 2 -# The 1..8-token MoE engine (router, routed experts at the TP8 width 384 and the shared expert in one launch) pushes -# its sum into the MoE boundary's all-reduce (prepushed mode 2). -MOE_ENGINE_MAX_TOKENS = 8 -MOE_PUSH_MAX_TOKENS = MOE_ENGINE_MAX_TOKENS -# 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. -MC_FINALIZE_MIN_TOKENS = MOE_ENGINE_MAX_TOKENS + 1 -# The router: the fp32 GEMV at one row, the tensor-core router (fused with the routed experts' NVFP4 activation -# quantization) at 2..48 rows. -ROUTER_GEMV_MAX_TOKENS = 1 -ROUTER_TC_MAX_TOKENS = 48 -# 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 (kernels/m3_shard_accept), at the configured draft length. The lm_head -# is the in-tree BF16 GEMM: at the TP8 shard (25008 rows) a coded copy of the head saves no time. -SHARD_ACCEPT_DRAFT_LEN = 3 -SHARD_ACCEPT_MAX_REQUESTS = DECODE_MAX_TOKENS // (SHARD_ACCEPT_DRAFT_LEN + 1) -# 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 (spec_worker). At TP8 it forks beside the verify head instead: forked early its -# GEMM lands beside the target tail it was meant to hide under. -FC_EARLY_MAX_ROWS = 0 -# 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 (kernels/m3_spec_tail.draft_input), the drafter layer's in-tree attention, MLP and -# all-reduces, and one kernel for the BF16 draft head, the cross-rank argmax and the pass's bookkeeping -# (kernels/m3_draft_step), at the configured draft length. The head kernel takes one row per request, at most 8, and -# streams the head's BF16 weight fragments. -DRAFT_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN -DRAFT_STEP_MAX_REQUESTS = DRAFT_STEP_KERNEL_MAX_TOKENS -# The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows -# (kernels/triton_sparse_decode.MQ_MIN_TOKENS). -SPARSE_DECODE_MQ_MIN_TOKENS = 17 -# Decode spans of at most this many requests score and select their top-16 blocks in one launch -# (ops/attention/m3_fused_score_select), which loads the older index-K pages before its grid wait; larger spans run -# the index scorer and the long-row selector. -FUSED_SCORE_SELECT_MAX_REQUESTS = FUSED_SCORE_SELECT_KERNEL_MAX_REQUESTS -# 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 -LONG_SELECT_MAX_BLOCKS = SELECT_LONG_MAX_BLOCKS -LONG_SELECT_KEY_WIDTHS = (2, 4, 9) -# The long-row select releases the sparse decode after it right after its own grid wait up to this many query rows; -# wider steps keep the trigger at the select's exit (the early-resident attention CTAs crowd the select grid). -SELECT_EARLY_MAX_ROWS = 8 -# The 8-head long-context sparse decode (ops/attention/m3_sparse_decode_lc8) stages the request length, its query rows -# and the request's block-table row before its grid wait up to this many query rows; above, the staged row's shared -# memory costs occupancy and the kernel stages only after the wait (same results bit for bit). -LC8_PREWAIT_MAX_ROWS = 8 -# The MSA index scorer (ops/attention/m3_index_score) prefetches the older index-K pages into L2 before its grid wait -# at every batch it runs (more than FUSED_SCORE_SELECT_MAX_REQUESTS requests); it never runs its one-request -# early-issue schedule here. -INDEX_SCORE_EARLY_MAX_REQUESTS = 0 - -# 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). -_ALL_ROWS = frozenset(range(1, DECODE_MAX_TOKENS + 1)) -MX_GEMV_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { - QKV_SPARSE: _ALL_ROWS, - QKV: _ALL_ROWS, - O_PROJ: _ALL_ROWS, - DENSE_GATE_UP: _ALL_ROWS, - DENSE_DOWN: _ALL_ROWS, -} -# 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: 48, - DENSE_GATE_UP: 16, - DENSE_DOWN: 32, -} -# The Eagle3 drafter's BF16 linears (qkv, o_proj, gate_up, down, fc) through FlashInfer's fixed-tactic low-M kernels -# (ops/gemm/m3_draft_lowm): per (N, K), the rows its tactic table covers; more rows keep the in-tree GEMM. -DRAFT_LOWM_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { - nk: frozenset(range(1, table[-1][0] + 1)) for nk, table in DRAFT_LOWM_TACTICS.items() -} - -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, -) -# The long-context sparse decode: "prewait" up to LC8_PREWAIT_MAX_ROWS query rows, "plain" above. -SPARSE_DECODE_LC8 = RowRange( - "m3_sparse_decode_lc8", - frozenset({"prewait", "plain"}), - 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. -PRODUCER = RowRange("m3_producer", frozenset({"nvfp4_tp8"}), m3_producer, _const("nvfp4_tp8")) -# Requests of a pure decode span whose blocks are scored and selected in one launch. -FUSED_SCORE_SELECT = RowRange( - "m3_fused_score_select", - _rows(1, FUSED_SCORE_SELECT_MAX_REQUESTS), - m3_fused_score_select, - _spec_rows("batch"), -) -INDEX_SCORE = RowRange( - "m3_index_score", - frozenset({"prefetch"}), - m3_index_score, - lambda spec: () if spec["early"] else ("prefetch",), -) -LONG_SELECT = RowRange( - "long_select", frozenset(LONG_SELECT_KEY_WIDTHS), m3_select_long, _long_select_widths -) -MC_ALLREDUCE = RowRange( - "mc_allreduce_norm", - _rows(1, DECODE_MAX_TOKENS), - 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), -) -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, DECODE_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") -) -ROUTER_GEMV = RowRange( - "m3_router_gemv", _rows(1, ROUTER_GEMV_MAX_TOKENS), m3_router_gemv, _spec_rows("m") -) -ROUTER_GEMV_FP4 = RowRange( - "m3_router_gemv_fp4", _rows(1, ROUTER_GEMV_MAX_TOKENS), m3_router_gemv_fp4, _spec_rows("m") -) -ROUTER_TC = RowRange( - "m3_router_tc", - _rows(ROUTER_GEMV_MAX_TOKENS + 1, ROUTER_TC_MAX_TOKENS), - m3_router_tc, - _spec_rows("m"), -) -# 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_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 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 (), -) - -# (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, - DENSE_DECODE_SPLIT, - SPARSE_DECODE_PRODUCER, - SPARSE_DECODE_LC8, - PRODUCER, - FUSED_SCORE_SELECT, - INDEX_SCORE, - LONG_SELECT, - MC_ALLREDUCE, - ATTN_PUSH, - MOE_PUSH, - MC_FINALIZE, - MOE_ENGINE, - ROUTER_GEMV, - ROUTER_GEMV_FP4, - ROUTER_TC, - MX_GEMV, - MX_GEMV_BF16, - DRAFT_LOWM, - SHARD_ACCEPT, - DRAFT_STEP, - 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 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 - - -def router_path(num_tokens: int) -> Optional[str]: - """The router kernel when the MoE runs unfused: the fp32 GEMV, the tensor-core router, or None (built-in gate).""" - if ROUTER_GEMV.covers(num_tokens): - return ROUTER_GEMV.path - if ROUTER_TC.covers(num_tokens): - return ROUTER_TC.path - return None - - -MX_GEMV_PATH = MX_GEMV.path -MX_GEMV_BF16_PATH = MX_GEMV_BF16.path -# 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) -> str: - """The decode-graph MXFP8 GEMM of an [N, K] projection on ``num_tokens`` rows. - - ``prequant``: the producer published the activation in MXFP8. - """ - 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 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() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/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_103__tp8/kernels/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_103__tp8/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/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_103__tp8/kernels/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_103__tp8/kernels/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/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_103__tp8/kernels/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_103__tp8/kernels/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/long_select.py index 4fb5064fa981..c01a82f0da41 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/long_select.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/long_select.py @@ -35,7 +35,7 @@ def long_select_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 (``dispatch.SELECT_EARLY_MAX_ROWS``). + grid wait (``SELECT_EARLY_MAX_ROWS`` in modeling.py). """ if ( max_score_kv.dtype != torch.float32 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_attn_lc8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_attn_lc8.py index 98b601471d10..2c4c55cf089f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_attn_lc8.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_attn_lc8.py @@ -15,9 +15,9 @@ 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 ``dispatch.LC8_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 (``dispatch.SELECT_EARLY_MAX_ROWS``). Neither changes a value. +weights while attention runs. Up to ``LC8_PREWAIT_MAX_ROWS`` (modeling.py) 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 (``SELECT_EARLY_MAX_ROWS`` in modeling.py). 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. @@ -25,10 +25,16 @@ import torch -from .. import dispatch from ..ops.attention.m3_sparse_decode_lc8 import NUM_HEADS, WIDTH, m3_sparse_decode_lc8 +def _prewait_max_rows() -> int: + """``LC8_PREWAIT_MAX_ROWS`` of the target's dispatch (imported at call time: modeling imports this module).""" + from .. import modeling + + return modeling.LC8_PREWAIT_MAX_ROWS + + 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: @@ -106,5 +112,5 @@ def m3_lc8_attn( float(sm_scale), bool(linear_softmax), bool(use_pdl), - total_q <= dispatch.LC8_PREWAIT_MAX_ROWS, + total_q <= _prewait_max_rows(), ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_draft_lowm.py deleted file mode 100644 index 356c26639947..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_draft_lowm.py +++ /dev/null @@ -1,121 +0,0 @@ -# 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 (``dispatch.DRAFT_LOWM``). - -``install`` swaps the linear method of every bias-free BF16 ``Linear`` of the drafter (qkv, o_proj, gate_up, down, the -fc; not the head) for ``DraftLowMLinearMethod``: a covered row count runs ``ops/gemm/m3_draft_lowm`` at the tactic -of its (N, K) table, every other call the in-tree method. ``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 Optional, Set, Tuple - -import torch -from torch import nn - -from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl -from tensorrt_llm._torch.modules.embedding import LMHead -from tensorrt_llm._torch.modules.linear import Linear, UnquantizedLinearMethod - -from .. import dispatch -from ..ops.gemm.m3_draft_lowm import m3_draft_lowm, tactic_for - -# (N, K, M, PDL) of every call prebuild compiled. -_BUILT: Set[Tuple[int, int, int, bool]] = set() - - -class DraftLowMLinearMethod(UnquantizedLinearMethod): - """The in-tree BF16 linear method with the drafter's covered rows on the fixed-tactic low-M kernels.""" - - def apply(self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor]): - out = draft_linear(module.weight, input, bias) - if out is not None: - return out - return super().apply(module, input, bias) - - -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 ``dispatch.DRAFT_LOWM`` covers it; None for the in-tree - GEMM.""" - if bias is not None or x.dim() != 2: - return None - n, k = weight.shape - m = int(x.shape[0]) - if not dispatch.DRAFT_LOWM.covers(((n, k), m)): - dispatch.uncovered( - dispatch.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. - dispatch.uncovered( - dispatch.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: - dispatch.missed(dispatch.DRAFT_LOWM.path, m, "the activation or weight is not BF16") - return None - if (n, k, m, use_pdl) not in _BUILT: - dispatch.missed( - dispatch.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) - out = torch.empty((m, n), dtype=x.dtype, device=x.device) - return m3_draft_lowm(x, weight, out, kind, tactic, use_pdl) - - -def _drafter_linears(draft_model: nn.Module): - for module in draft_model.modules(): - if ( - isinstance(module, Linear) - and not isinstance(module, LMHead) - and type(module.quant_method) is UnquantizedLinearMethod - and getattr(module, "bias", None) is None - and module.weight.dtype == torch.bfloat16 - ): - yield module - - -def install(draft_model: nn.Module) -> int: - """Put every covered drafter linear on ``DraftLowMLinearMethod``; returns how many were installed.""" - count = 0 - for module in _drafter_linears(draft_model): - n, k = module.weight.shape - if any( - dispatch.DRAFT_LOWM.covers(((n, k), m)) - for m in range(1, dispatch.DECODE_MAX_TOKENS + 1) - ): - module.quant_method = DraftLowMLinearMethod() - count += 1 - return count - - -def prebuild(draft_model: nn.Module) -> int: - """Compile (and run once) every covered row count of every installed drafter linear, outside capture, on the - current stream. Returns the number of calls compiled.""" - use_pdl = get_env_enable_pdl() - built = 0 - shapes = { - tuple(module.weight.shape): module.weight - for module in draft_model.modules() - if isinstance(module, Linear) and isinstance(module.quant_method, DraftLowMLinearMethod) - } - for (n, k), weight in sorted(shapes.items()): - for nk, m in sorted(key for key in dispatch.DRAFT_LOWM.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) - _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_103__tp8/kernels/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_draft_step.py index 42aee2cb9041..917f04005cc9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_draft_step.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_draft_step.py @@ -25,7 +25,6 @@ import torch -from .. import dispatch 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 .m3_shard_accept import AcceptWorkspace, create_workspace @@ -79,13 +78,13 @@ def _num_sms() -> int: 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.""" + """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(): - dispatch.capture_missed(dispatch.DRAFT_STEP.path, None, "the draft-step exchange workspace") + # 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 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_force_vec.py new file mode 100644 index 000000000000..b4de37794ca9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_force_vec.py @@ -0,0 +1,35 @@ +# 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``), drawn exactly as its ``_apply_force_accepted_tokens`` does, in one launch. + +Triton: 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. Reached through the ``ops/spec`` entry ``m3_force_vec``. +""" + +import triton +import triton.language as tl + + +@triton.jit +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) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_index_score.py index 9748bd1fe4a5..dfe1909e537a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_index_score.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_index_score.py @@ -543,10 +543,10 @@ class IndexScoreRunner: per CTA.""" kernel_cache: dict = {} - SUPPORTED_HEAD_DIM = 128 + 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 + MAX_BLOCK_Q = 32 # capacity: fixed: the scorer's query tile EARLY_PREFETCH_DEPTH = 1 @classmethod diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_mc_allreduce.py index 7617e9f64772..823c78a4ba60 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_mc_allreduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_mc_allreduce.py @@ -38,7 +38,7 @@ m3_mc_finalize_allreduce_norm as _finalize_allreduce_norm_op, ) -MAX_TOKENS = 48 +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 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/m3_moe_kernel.py index a45f0cee2ea3..34592d700b99 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/m3_moe_kernel.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/m3_moe_kernel.py @@ -62,7 +62,7 @@ def _cfg(key: str, default): # ============================================================================= # Problem shape (trace-time constants). # ============================================================================= -H = 6144 +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) @@ -70,7 +70,7 @@ def _cfg(key: str, default): TWO_I = 2 * I_TP assert I_TP % 128 == 0 and 128 <= I_TP <= 768, I_TP E_LOCAL = 128 -TOP_K = 4 +TOP_K = 4 # capacity: fixed: the checkpoint's top-k ROUTE_SCALE = 2.0 SUM_EPS = 1e-20 M_MAX = _cfg("m_max", 1) @@ -91,7 +91,7 @@ def _cfg(key: str, default): 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 +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 @@ -119,12 +119,14 @@ def _cfg(key: str, default): # 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 # gate_up K blocks -SH_DKB = SH_I // 32 # down K blocks +SH_KB = H // 32 # gate_up K blocks # capacity: fixed: a kernel tile constant +SH_DKB = SH_I // 32 # down K blocks # capacity: fixed: a kernel tile constant SH_DKT = SH_I // 128 # down k-tiles -SH_GT = SH_KB // 4 // SH_SPLIT # gate_up k-tiles (128 K) per slice +# Gate_up k-tiles (128 K) per slice. +SH_GT = SH_KB // 4 // SH_SPLIT # capacity: fixed: a kernel tile constant SH_STAGES_GU = SH_GT -SH_STAGES = SH_GT + 2 * SH_DKT # + 2 down tiles x SH_DKT k-tiles +# + 2 down tiles x SH_DKT k-tiles. +SH_STAGES = SH_GT + 2 * SH_DKT # capacity: fixed: pipeline stages 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 @@ -149,14 +151,14 @@ def _cfg(key: str, default): 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 +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 +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 @@ -186,7 +188,7 @@ def _nkb2(t2: int) -> int: 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 +PARTIAL_EMPTY = -1 # capacity: fixed: a sentinel, not a size EPI_THREADS = 128 EPI_BAR = 1 ROUTE_BAR = 2 @@ -229,11 +231,12 @@ def _nkb2(t2: int) -> int: 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) +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)) @@ -308,8 +311,8 @@ def _nkb2(t2: int) -> int: 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 +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 FUSED_ROUTER, "shared expert: fused router only" assert SH_STAGES - SH_NST <= 64, SH_STAGES diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/op.py index d9d3fe220d8d..79d4e637c251 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/op.py @@ -53,19 +53,23 @@ import torch -HIDDEN = 6144 -INTER = 384 # the routed intermediate per TP rank +HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size +# The routed intermediate per TP rank. +INTER = 384 # capacity: fixed: the routed experts' width per tp8 rank NUM_EXPERTS = 128 -SHARED_INTER = 384 # the shared expert's intermediate per TP rank (the routed one) -SHARED_CTAS = 24 # m3_moe's shared-expert CTAs (beside num_ctas routed CTAs) +# The shared expert's intermediate per TP rank (the routed one). +SHARED_INTER = 384 # capacity: fixed: the shared expert's width per tp8 rank +# M3_moe's shared-expert CTAs (beside num_ctas routed CTAs). +SHARED_CTAS = 24 # capacity: fixed: an SM split, not a size 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) +# 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 -_ITEM_BYTES = _UNIT_BYTES // _S1 +_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). @@ -84,7 +88,7 @@ # ``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 +LG_REPL_MAX_TOKENS = 4 # capacity: fixed: measured exchange-replica 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 @@ -198,7 +202,9 @@ def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast 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) + 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 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_router_gemv.py index 88a475f509d5..30855eb7f05d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_router_gemv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_router_gemv.py @@ -20,7 +20,7 @@ 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 +_MAX_DYNAMIC_SMEM = 227 * 1024 - 1024 # capacity: fixed: the device's shared-memory limit def m3_router_gemv( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_sampled_accept.py index 315e80947213..59c58336c5d7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_sampled_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_sampled_accept.py @@ -28,11 +28,12 @@ import torch -from .. import dispatch +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/minimaxM3/sampledAccept.cu) and called as -# torch.ops.trtllm.minimax_m3_sampled_accept. -MAX_ROWS = 48 +# The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.cu) and reached through +# the ops/spec entry m3_sampled_accept. +MAX_ROWS = 48 # capacity: fixed: the kernel's built rows (decode steps) class SampledAcceptWorkspace: @@ -49,7 +50,7 @@ def __init__(self, uc, mc_ptr, rank, world, handle=None, comm=None): @classmethod def create(cls, mapping, device) -> "SampledAcceptWorkspace": - words = int(torch.ops.trtllm.minimax_m3_sampled_accept_ws_words(max(mapping.tp_size, 1))) + words = int(_ws_words_op(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) @@ -88,7 +89,7 @@ def workspace(mapping, device) -> Optional[SampledAcceptWorkspace]: global _WS if _WS is None: if torch.cuda.is_current_stream_capturing(): - dispatch.capture_missed("spec_tail", None, "the sampled-accept workspace") + # 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) return _WS @@ -131,7 +132,7 @@ def sampled_accept( 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_sampled_accept( + _sampled_accept_op( logits, int(vocab_offset), temperatures[:rows].float().contiguous(), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_shard_accept.py index e1dafd292d6d..2b7d05ad920f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_shard_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_shard_accept.py @@ -13,7 +13,7 @@ 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``. +shardAccept.cu) and reached through the ``ops/spec`` entry ``m3_shard_accept``. """ from dataclasses import dataclass @@ -21,10 +21,10 @@ import torch -from .. import dispatch +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 +SHARD_ROWS = 48 # capacity: fixed: DECODE_MAX_TOKENS verify rows @dataclass(eq=False) @@ -107,7 +107,7 @@ def _workspace(mapping, device) -> Optional[AcceptWorkspace]: if _WS is not None: return _WS if torch.cuda.is_current_stream_capturing(): - dispatch.capture_missed(dispatch.SHARD_ACCEPT.path, None, "the shard-accept workspace") + # 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 @@ -154,7 +154,7 @@ def shard_accept( 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( + _shard_accept_op( logits, vocab_offset, ws.keys, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_spec_tail.py index 996371370fde..0d4dc9cd1b64 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_spec_tail.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_spec_tail.py @@ -13,8 +13,13 @@ 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}. +# 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( @@ -33,10 +38,9 @@ def draft_input( 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( + return _draft_input_op( ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, use_pdl ) - return out, residual def chunk_norm( @@ -45,8 +49,7 @@ def chunk_norm( """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 + 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 / @@ -72,7 +75,7 @@ def step_gather( and kv_off.is_contiguous() ): return False - torch.ops.trtllm.minimax_m3_step_gather( + _step_gather_op( new_tok, next_draft, lens, slots, ppi, P, T, D, ids, ids0, draft, d0, pos_off, kv_off, b0 ) return True @@ -88,7 +91,7 @@ def slot_scatter( 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( + _slot_scatter_op( slots, src_new, dst_new, @@ -110,4 +113,4 @@ def capture_norm( 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) + _capture_norm_op(h, residual, weight, eps, dst) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/m3_mx_gemv_kernel.py index b1c9d4de0e08..e3d7751231ee 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/m3_mx_gemv_kernel.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/m3_mx_gemv_kernel.py @@ -40,11 +40,13 @@ from cutlass import dsl_user_op from cutlass.experimental import primitives as prims -CTA_M = 128 # weight rows per CTA = the MMA's M +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 +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 +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 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/op.py index 61c5f6962e2f..84808d7c61bf 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/op.py @@ -28,14 +28,14 @@ from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl -MAX_TOKENS = 48 +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 +_SMEM_LIMIT = 220 * 1024 # capacity: fixed: the device's shared-memory limit def _mma_n(m: int) -> int: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/triton_sparse_decode.py index fde91fd41c5a..0e8dc710bc98 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/triton_sparse_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/triton_sparse_decode.py @@ -57,7 +57,7 @@ from . import m3_attn_lc8, m3_prequant # One sparse block is exactly one KV page. -SPARSE_BLOCK_SIZE = 128 +SPARSE_BLOCK_SIZE = 128 # capacity: fixed: the MSA block size of the checkpoint _FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) @@ -68,7 +68,7 @@ } # Total CTAs the split-K partitioning aims for before it stops splitting. -_TARGET_GRID = 256 +_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. @@ -77,7 +77,7 @@ # 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 +MERGE_MXFP8_MIN_TOKENS = 1 # capacity: fixed: a row lower bound class MergeMxfp8: @@ -89,7 +89,7 @@ class MergeMxfp8: merge rewrites it, in stream order. """ - MAX_ROWS = 128 + MAX_ROWS = 128 # capacity: fixed: the merge hand-off's rows (decode steps) def __init__(self): self._buffers = {} @@ -1176,7 +1176,7 @@ def _check_nvfp4_inputs( # 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 +MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query producer's crossover # Warps per multi-query producer program. MQ_NUM_WARPS = 8 @@ -1333,7 +1333,7 @@ def minimax_m3_sparse_attn_decode( 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 is the kernel's (dispatch.SPARSE_DECODE_LC8); operands it declines are a miss. + # cache with FP8 q at the TP8 heads is the kernel's (modeling's SPARSE_DECODE_LC8); operands it declines are a miss. lc8_covered = ( kv_nvfp4 and widen_q @@ -1345,10 +1345,10 @@ def minimax_m3_sparse_attn_decode( if lc8_covered and not m3_attn_lc8.lc8_applies( q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output ): - from .. import dispatch + from .. import modeling - dispatch.missed( - dispatch.SPARSE_DECODE_LC8.path, + modeling.missed( + modeling.SPARSE_DECODE_LC8.path, total_q, "the step's q / KV pages / scales / output are not the layout the kernel is built for", ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/layers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/layers.py deleted file mode 100644 index f0bdf31bb4c7..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/layers.py +++ /dev/null @@ -1,1403 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""MiniMax-M3 NVFP4 / sm_103 / tp8: the decoder layers, on the built-in modules. - -The attention and MoE modules subclass the built-in ``MiniMaxM3Attention`` and -``MiniMaxM3MoE``: the generic path (prefill, mixed steps, wide steps) runs their -code unchanged, and the checkpoint loads through the built-in loader. What this -target adds is the decode-step kernels, each taken only by the step shapes it -was written for and measured on (GB300, tp8 over two trays, 1-8 requests): - -* attention: the NVFP4 QKV + indexer producer at the 8-head geometry; for a - long-context generation step of up to 2 requests the index scoring and the - top-16 block selection in one launch, else the index scorer and the long-row - selector; the selected blocks' attention, split-K merge and the o_proj's MXFP8 - hand-off in one 8-head launch; -* MoE: the router (an fp32 GEMV or a tensor-core router fused with the routed - experts' NVFP4 activation quantization) and the routed experts (width 384) as - one engine launch for 1-8 tokens, with the shared expert folded in; -* the layer boundaries: a multicast all-reduce fused with the residual add and - the next RMSNorm up to 48 tokens. - -The decoder layer and the text model are this target's own: they construct the -subclassed modules. Their generic forward is the built-in one. -""" - -from __future__ import annotations - -from typing import Callable, Dict, Optional, Tuple - -import torch -from torch import nn -from transformers import PretrainedConfig - -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.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention -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_minimaxm3 import ( - MiniMaxM3Attention, - MiniMaxM3MoE, - _build_swiglu_oai_dense_mlp, - _nvfp4_cache_supports_fused_write, - _validate_sparse_attention_runtime_config, - get_moe_layer_ids, - get_sparse_disable_index_value_layer_ids, - get_sparse_layer_ids, - maybe_bcg_minimax_m3_attn_custom_op_inplace, -) -from tensorrt_llm._torch.models.modeling_utils import DecoderModel -from tensorrt_llm._torch.modules.decoder_layer import DecoderLayer -from tensorrt_llm._torch.modules.embedding import Embedding -from tensorrt_llm._torch.modules.linear import flashinfer_mxfp8_autotune_active -from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel -from tensorrt_llm._torch.modules.rms_norm import RMSNorm -from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph -from tensorrt_llm._torch.speculative import SpecMetadata -from tensorrt_llm._torch.utils import ( - AuxStreamType, - EventType, - Fp4QuantizedTensor, - is_torch_compiling, -) -from tensorrt_llm.logger import logger - -from . import dispatch -from .dispatch import MOE_ENGINE_MAX_TOKENS, PRODUCER_EARLY_TRIGGER, ROUTER_TC_MAX_TOKENS -from .kernels import m3_mc_allreduce, m3_prequant -from .kernels.m3_router_gemv import m3_router_gemv, m3_router_gemv_fp4, m3_router_tc -from .linears import PUSH_ATTN, PUSH_MOE, mark_m3_linear -from .ops.attention.m3_index_score import SPLIT_K, m3_index_score -from .ops.attention.m3_producer import m3_producer -from .ops.moe.m3_moe_engine import m3_moe_engine -from .spec_worker import capture_hidden_states - -# The row ranges of the engines, the router and the fused attention are in dispatch.py. -# The shared expert folds into the routed engine from one token (m3_moe with dedicated shared-expert CTAs). -MOE_FOLD_MIN_TOKENS = 1 -# SMs left to the concurrent shared-expert kernels while a routed engine without the shared fold runs (persistent). -MOE_ENGINE_RESERVED_SMS = 24 -# CTAs per request of the CuTe DSL index scorer at every batch (each block is scored once, so the scores do not depend -# on it), and the request counts up to which it compiles the early-issue schedule (kernels/m3_index_score). -INDEX_SCORE_SPLIT_K = SPLIT_K -INDEX_SCORE_EARLY_MAX_REQUESTS = dispatch.INDEX_SCORE_EARLY_MAX_REQUESTS -# The generic layer's boundaries (prefill, Eagle3 capture layers): the low-latency multicast all-reduce + residual + -# RMSNorm beats the native one-shot all-reduce up to 8 tokens; with the shared + routed add folded in, up to its own -# maximum. -MC_AR_MAX_TOKENS = 8 - -# 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] = {} -# The CuTe DSL index scorer of the sparse decode (its kernels are compiled by prebuild_index_scorer). -INDEX_SCORE_PATH = dispatch.INDEX_SCORE.path - - -def _scorer_runner(): - from .kernels.m3_index_score import IndexScoreRunner - - return IndexScoreRunner - - -def _scorer_key( - dtype: torch.dtype, num_heads: int, query_len: int, head_dim: int, split_k: int, early: bool -) -> tuple: - """The scorer's compile key (``IndexScoreRunner``).""" - return (dtype, int(num_heads), int(query_len), int(head_dim), int(split_k), bool(early)) - - -def scorer_variants(layers, query_lens) -> list: - """The scorer keys a decode step of these layers can take: every sparse layer's (index dtype, index heads per - rank, index dim) at each decode query length, with the early-issue schedule (up to - ``INDEX_SCORE_EARLY_MAX_REQUESTS`` requests, none here) and without it, where the kernel serves the geometry.""" - runner = _scorer_runner() - keys = [] - for layer in layers: - attn = getattr(layer, "self_attn", None) - if ( - attn is None - or not getattr(attn, "is_sparse_attention_layer", False) - or not isinstance(attn.attn, MiniMaxM3MsaSparseAttention) - ): - continue - config = attn.attn.m3_config - dtype = torch.float8_e4m3fn if attn.attn.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) if INDEX_SCORE_EARLY_MAX_REQUESTS else (False,): - key = _scorer_key( - dtype, - config.num_index_heads, - qlen, - config.sparse_index_dim, - INDEX_SCORE_SPLIT_K, - early, - ) - if key not in keys: - keys.append(key) - return keys - - -def prebuild_index_scorer(layers, query_lens, device) -> list: - """Compile every key of ``scorer_variants`` 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 = _scorer_runner() - keys = scorer_variants(layers, 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 - - -_EXPERT_TENSOR_NAMES = ( - "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", -) - - -def _nvfp4_swiglu_oai_experts(experts: nn.Module): - """The routed experts' operands the engines read, or raise if the backend is not TRTLLM-Gen NVFP4 SwiGLU-OAI.""" - backend = getattr(experts, "backend", experts) - tensors = [getattr(backend, name, None) for name in _EXPERT_TENSOR_NAMES] - scale = getattr(backend, "fc31_input_scale", None) - if ( - any(t is None for t in tensors) - or scale is None - or getattr(backend, "scaling_vector_size", 16) != 16 - or getattr(backend, "fc31_act_scale", None) is not None - or not type(backend).quantize_input.__qualname__.startswith("TRTLLMGenNvfp4Base") - ): - raise ValueError("experts backend is not TRTLLM-Gen NVFP4 SwiGLU-OAI") - tensors = [t.data if isinstance(t, nn.Parameter) else t for t in tensors] - scale = scale.data if isinstance(scale, nn.Parameter) else scale - return tensors, scale - - -def _dense_mlp(model_config: ModelConfig, intermediate_size: int, layer_idx: int) -> nn.Module: - """The built-in dense SwiGLU-OAI MLP, its projections on this target's decode GEMM dispatch.""" - mlp = _build_swiglu_oai_dense_mlp( - model_config=model_config, intermediate_size=intermediate_size, layer_idx=layer_idx - ) - mark_m3_linear(mlp.gate_up_proj) - mark_m3_linear(mlp.down_proj) - return mlp - - -class M3MoE(MiniMaxM3MoE): - """MiniMax-M3 MoE with the decode-step router and expert engines. - - Engines are built per (token count, fusion) by ``prebuild_engines`` from - ``build_targets``, before the executor's warmup (building compiles and warms - the kernel); a forward only looks them up, and a step that finds no engine - is a miss and runs the built-in path. - """ - - def __init__( - self, - model_config: ModelConfig, - aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - layer_idx: int, - ) -> None: - super().__init__(model_config, aux_stream_dict, layer_idx) - self.layer_idx = layer_idx - if self.shared_experts is not None: - mark_m3_linear(self.shared_experts.gate_up_proj) - mark_m3_linear(self.shared_experts.down_proj, push=PUSH_MOE) - # The routed experts' NVFP4 input scale once resolved; False when they cannot take the fused router. - self._fp4_input_scale: Optional[torch.Tensor | bool] = None - self._engines: Dict[Tuple[int, bool, bool], Optional[object]] = {} - - def _router_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: - """The fp32 router logits: the fp32 GEMV at one row, the tensor-core router up to 48.""" - rows = hidden_states.shape[0] if hidden_states.dim() == 2 else 0 - path = dispatch.router_path(rows) - if path == dispatch.ROUTER_GEMV.path: - logits = m3_router_gemv(hidden_states, self.gate.weight, use_pdl=get_env_enable_pdl()) - if logits is not None: - return logits - dispatch.missed(path, rows, "the router GEMV declined the operands") - elif path == dispatch.ROUTER_TC.path: - fused = m3_router_tc(hidden_states, self.gate.weight, use_pdl=get_env_enable_pdl()) - if fused is not None: - return fused[0] - dispatch.missed(path, rows, "the tensor-core router declined the operands") - else: - dispatch.uncovered(dispatch.ROUTER_TC.path, f"more than {ROUTER_TC_MAX_TOKENS} rows") - return self.gate(hidden_states) - - def _router_fp4(self, hidden_states: torch.Tensor): - """Router logits and the routed experts' NVFP4 activation from one kernel, or None (then - ``_router_logits``).""" - if hidden_states.dim() != 2 or hidden_states.shape[0] > ROUTER_TC_MAX_TOKENS: - return None - if self._fp4_input_scale is None: - try: - _, scale = _nvfp4_swiglu_oai_experts(self.experts) - weight = getattr(self.experts, "backend", self.experts).w3_w1_weight - self._fp4_input_scale = ( - scale if weight.shape[-1] * 2 == hidden_states.shape[-1] else False - ) - except ValueError: - self._fp4_input_scale = False - logger.info( - f"[M3 router FP4] fused router+NVFP4 activation quant: {self._fp4_input_scale is not False}" - ) - rows = hidden_states.shape[0] - path = dispatch.router_path(rows) - if path == dispatch.ROUTER_GEMV.path: - path = dispatch.ROUTER_GEMV_FP4.path - if self._fp4_input_scale is False: - dispatch.missed( - path, - rows, - "the routed experts are not TRTLLM-Gen NVFP4 SwiGLU-OAI on this hidden size", - ) - return None - if path == dispatch.ROUTER_TC.path: - fused = m3_router_tc( - hidden_states, self.gate.weight, self._fp4_input_scale, use_pdl=get_env_enable_pdl() - ) - else: - fused = m3_router_gemv_fp4( - hidden_states, self.gate.weight, self._fp4_input_scale, use_pdl=get_env_enable_pdl() - ) - if fused is None: - dispatch.missed(path, rows, "the router declined the operands") - return fused - - def engine(self, num_tokens: int, shared: bool) -> Optional[object]: - """This layer's prebuilt routed-expert engine for ``num_tokens`` (1..8), router folded in. - - ``shared`` folds the shared expert in as well (routed + shared out). None - when the built-in path must run, after a miss: the engine failed to build, - or ``prebuild_engines`` did not build it (a forward never builds one). - """ - key = (num_tokens, True, shared) - return self._lookup(self._engines, key, dispatch.MOE_ENGINE.path, num_tokens) - - def _lookup(self, engines: dict, key, path: str, num_tokens: int) -> Optional[object]: - if key not in engines: - what = f"layer {self.layer_idx}'s engine" - if torch.cuda.is_current_stream_capturing(): - dispatch.capture_missed(path, num_tokens, what) - else: - dispatch.missed(path, num_tokens, f"{what} was not built before warmup") - return None - engine = engines[key] - if engine is None: - dispatch.missed(path, num_tokens, "its build failed (logged)") - return engine - - def prebuild_engines(self) -> None: - """Build (compile, on first use of each shared state, and warm) every engine a decode step of this layer can - take: ``MOE_ENGINE``'s token counts with the shared expert folded in. Runs from - ``build_targets`` once the weights are real, on the current stream, in the same order on every rank; a build - that fails is logged and its token count falls back (a miss) in the forward. Whether an engine is kept is - agreed across the TP group (one host collective per engine, here): a build that fails on one rank only drops - that engine on every rank, so the ranks never mix an engine (and its boundary push) with the built-in MoE.""" - if self.shared_experts is None: - return - for num_tokens in sorted(dispatch.MOE_ENGINE.keys): - key = (num_tokens, True, True) - if key not in self._engines: - self._engines[key] = self._agreed( - dispatch.MOE_ENGINE.path, - num_tokens, - self._try_build( - dispatch.MOE_ENGINE.path, num_tokens, self._build_engine, shared=True - ), - ) - - def _agreed(self, path: str, num_tokens: int, engine: Optional[object]) -> Optional[object]: - return m3_mc_allreduce.tp_agree_built( - self.mapping, - engine, - f"[M3 dispatch] {path} at {num_tokens} rows: layer {self.layer_idx}'s engine", - ) - - def _try_build(self, path: str, num_tokens: int, build, **kwargs) -> Optional[object]: - try: - return build(num_tokens, **kwargs) - except (ImportError, RuntimeError, ValueError) as err: - logger.warning( - f"[M3 dispatch] {path} at {num_tokens} rows: layer {self.layer_idx}'s engine failed to build: {err}" - ) - return None - - def _build_engine(self, num_tokens: int, shared: bool) -> object: - from .ops.moe import m3_moe_engine as m3_moe_op - - tensors, scale = _nvfp4_swiglu_oai_experts(self.experts) - device = tensors[0].device - sms = torch.cuda.get_device_properties(device).multi_processor_count - kwargs = {} - if shared: - gu, dn = self.shared_experts.gate_up_proj, self.shared_experts.down_proj - kwargs["shared_weights"] = ( - gu.weight.data.view(torch.float8_e4m3fn), - gu.weight_scale.data, - dn.weight.data.view(torch.float8_e4m3fn), - dn.weight_scale.data, - ) - state = _ENGINE_STATES.get((str(device), num_tokens, shared)) - if state is None: - state = _ENGINE_STATES[(str(device), num_tokens, shared)] = m3_moe_op.M3MoeState( - device, - num_tokens=num_tokens, - num_ctas=sms - MOE_ENGINE_RESERVED_SMS, - shared=shared, - ) - engine = state.layer( - *tensors, - self.gate.e_score_correction_bias.data, - router_weight=self.gate.weight.data, - fc31_input_scale=scale.reshape(1).float(), - **kwargs, - ) - # Compile (first use of the state) and warm. - engine( - torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device) - ) - return engine - - def forward( - self, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - final_all_reduce_params: Optional[AllReduceParams] = None, - lora_params: Optional[dict] = None, - do_finalize: bool = True, - return_parts: bool = False, - ): - """The MoE output; or, with ``return_parts``, ``(shared, routed)`` (routed None when the - engine folded the shared expert in); or, with ``do_finalize=False``, the unreduced - ``[shared_output, fc2_output, expert_scale_factor, expanded_idx_to_permuted_idx]`` - for the fused finalize + all-reduce + residual + RMSNorm epilogue. - - Mirrors the built-in ``MiniMaxM3MoE.forward`` apart from the decode engines. - """ - all_rank_num_tokens = attn_metadata.all_rank_num_tokens - engine_rows = ( - do_finalize - and hidden_states.dim() == 2 - and hidden_states.dtype == torch.bfloat16 - and 1 <= hidden_states.shape[0] <= MOE_ENGINE_MAX_TOKENS - ) - - def _compute_routed_output(): - fused = self._router_fp4(hidden_states) - if fused is not None: - router_logits, x_fp4, x_sf = fused - return self.experts( - Fp4QuantizedTensor(x_fp4, x_sf, is_sf_swizzled=False), - router_logits, - do_finalize=do_finalize, - output_dtype=hidden_states.dtype, - all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=False, - lora_params=lora_params, - ) - router_logits = self._router_logits(hidden_states) - return self.experts( - hidden_states, - router_logits, - do_finalize=do_finalize, - all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=False, - lora_params=lora_params, - ) - - def _compute_shared_output(): - return self.shared_experts(hidden_states, lora_params=lora_params) - - folded = None - if ( - engine_rows - and self.shared_experts is not None - and hidden_states.shape[0] >= MOE_FOLD_MIN_TOKENS - ): - engine = self.engine(hidden_states.shape[0], shared=True) - if engine is not None: - arp = m3_prequant.MOE_PUSH.take_engine() - folded = m3_moe_engine(hidden_states.contiguous(), engine, arp=arp) - if arp is not None: - m3_prequant.MOE_PUSH.engine_pushed = folded.data_ptr() - if folded is not None: - # Routed + shared from one kernel. - if return_parts: - return folded, None - result = folded - elif self.shared_experts is None: - result = _compute_routed_output() - else: - routed_output, shared_output = maybe_execute_in_parallel( - _compute_routed_output, - _compute_shared_output, - self.event_dict[EventType.Main], - self.event_dict[EventType.MoeShared], - # 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 self.aux_stream, - disable_on_compile=True, - ) - if not do_finalize: - return [shared_output, *routed_output] - if return_parts: - return shared_output, routed_output - # In-place add into ``shared_output`` to avoid allocating a temporary. - result = shared_output.add_(routed_output) - # tp8 routed experts are tensor-parallel (never fused-communication), so the sum is still rank-local. - return self.allreduce(result, all_reduce_params=final_all_reduce_params) - - -def producer_pw_covered(layer_uses_nvfp4, num_heads_q, num_heads_kv, num_heads_index) -> bool: - """Whether ``dispatch.PRODUCER`` covers a sparse layer: 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 bool(layer_uses_nvfp4) and (num_heads_q, num_heads_kv, num_heads_index) == (8, 1, 1) - - -def _producer_miss(covered: bool, rows: int, reason: str) -> None: - """A producer precondition failed: a miss of ``dispatch.PRODUCER`` on a covered layer (raised under - ``modeling_v2='require'``), otherwise nothing; the caller then runs the unfused producer.""" - if covered: - dispatch.missed(dispatch.PRODUCER.path, rows, reason) - return None - - -def producer_pw_supported(num_heads_q, num_heads_kv, num_heads_index, kv_data, kv_scale) -> bool: - """Whether a sparse layer's NVFP4 producer can take ``m3_producer``: the TP8 heads (8 query, one KV and one - index head) and the NVFP4 pool [pages, 2, 1, 128, 64] with its scale pool [pages, 2, 1, 128, 8].""" - return ( - num_heads_q == 8 - and num_heads_kv == 1 - and num_heads_index == 1 - and 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] - ) - - -class M3Attention(MiniMaxM3Attention): - """MiniMax-M3 attention with the decode-step fused kernels.""" - - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - mark_m3_linear(self.qkv_proj) - mark_m3_linear(self.o_proj, push=PUSH_ATTN) - if not self.is_sparse_attention_layer: - # The KV split trtllm-gen dense decode runs at, per request count (msa_decode.dense_decode_max_seq_len). - self.attn.dense_decode_max_seq_len_by_batch = dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH - - def _check_producer_projection(self) -> None: - """A layer the producer covers (dispatch.PRODUCER) needs the fused QKV/index projection: without it the step is - a miss (raises under require) instead of silently taking the unfused path.""" - if not self.enable_fused_qkv_index_projection and producer_pw_covered( - self.main_kv_is_nvfp4, - self.num_heads, - self.num_key_value_heads, - self.sparse_num_index_heads, - ): - dispatch.missed(dispatch.PRODUCER.path, None, "fuse_qkv_index_projection is disabled") - - def _fused_fp8_qkv_indexer_norm_rope_kv_insert( - self, - packed: torch.Tensor, - position_ids: Optional[torch.Tensor], - attn_metadata: AttentionMetadata, - *, - cache_tensors: Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = None, - ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: - """Run the vLLM-style horizontal producer for every sparse batch. - - Mirrors the built-in ``MiniMaxM3Attention._fused_fp8_qkv_indexer_norm_rope_kv_insert``; an NVFP4 cache takes - this target's ``m3_producer`` (the native producer bit for bit, its constant loads staged before the grid - wait and its dependents triggered on entry) where its geometry is the one that kernel is built for. - - The CUDA kernel is token-major and batch-type agnostic: per-token - positions and cache slots cover pure prefill, mixed aggregate batches, - and CUDA-graph decode. FP8 caches receive FP8 K/V directly. NVFP4 - caches receive packed E2M1 K/V plus their E4M3 SF16 scale bytes while - Q, index-Q, and index-K remain FP8. - """ - if ( - not self.enable_fused_qkv_index_projection - or not isinstance(self.attn, MiniMaxM3MsaSparseAttention) - or not (self.main_kv_is_fp8 or self.main_kv_is_nvfp4) - or self.attn.indexer_kv_dtype != "fp8" - ): - return None - # A sparse layer on an NVFP4 cache at the TP8 heads is the producer's (dispatch.PRODUCER): every precondition - # below that fails for it is a miss. - covered = producer_pw_covered( - self.main_kv_is_nvfp4, - self.num_heads, - self.num_key_value_heads, - self.sparse_num_index_heads, - ) - rows = int(packed.shape[0]) - rope = self.pos_embd_params.rope if self.pos_embd_params is not None else None - rotary_dim = int(rope.dim) if rope is not None else 0 - if ( - packed.dtype != torch.bfloat16 - or position_ids is None - or self.head_dim != 128 - or self.sparse_index_dim != 128 - or self.sparse_num_index_heads != self.num_key_value_heads - or not self.use_gemma_norm - or self.pos_embd_params is None - or not self.pos_embd_params.is_neox - or self.rotary_emb is None - or rope is None - or rotary_dim != 64 - ): - return _producer_miss( - covered, - rows, - "the packed QKV dtype, position ids, head dims, norm or RoPE configuration", - ) - norm_eps = self.q_norm.variance_epsilon - if any( - module.variance_epsilon != norm_eps - for module in (self.k_norm, self.index_q_norm, self.index_k_norm) - ): - return _producer_miss(covered, rows, "the four RMSNorm epsilons differ") - norm_weights = ( - self.q_norm.weight, - self.k_norm.weight, - self.index_q_norm.weight, - self.index_k_norm.weight, - ) - if any(weight.dtype != torch.bfloat16 or not weight.is_cuda for weight in norm_weights): - return _producer_miss(covered, rows, "a norm weight is not BF16 on the GPU") - - kv_cache_manager = attn_metadata.kv_cache_manager - if cache_tensors is None: - if kv_cache_manager is None: - return _producer_miss(covered, rows, "the step has no KV cache manager") - buffers = kv_cache_manager.get_buffers(self.layer_idx, kv_layout="HND") - index_k_cache = attn_metadata.msa_idx_k_cache(self.layer_idx) - out_cache_loc = getattr(attn_metadata, "msa_out_cache_loc", None) - else: - buffers, index_k_cache, out_cache_loc = cache_tensors - num_tokens = int(packed.shape[0]) - layer_uses_nvfp4 = bool( - getattr(kv_cache_manager, "is_nvfp4_layer", lambda _layer_idx: False)(self.layer_idx) - ) - scale_buffers = ( - kv_cache_manager.get_block_scale_buffers(self.layer_idx, kv_layout="HND") - if layer_uses_nvfp4 - else None - ) - kv_quant_scale = getattr(self.qkv_proj, "inv_kv_scales", None) if layer_uses_nvfp4 else None - if layer_uses_nvfp4: - supports_fused_layout = ( - buffers is not None - and scale_buffers is not None - and _nvfp4_cache_supports_fused_write( - buffers, scale_buffers, self.num_key_value_heads - ) - ) - supported_main_cache = ( - supports_fused_layout - and buffers.is_cuda - and scale_buffers.is_cuda - and scale_buffers.device == buffers.device - and kv_quant_scale is not None - 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_cache = ( - buffers is not None - and buffers.is_cuda - and buffers.dtype == torch.float8_e4m3fn - and buffers.dim() == 5 - and tuple(buffers.shape[1:]) == (2, self.num_key_value_heads, 128, 128) - and buffers.stride(4) == 1 - and buffers.stride(3) == 128 - and buffers.stride(2) == 128 * 128 - and buffers.stride(1) >= self.num_key_value_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_cache = ( - 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] - ) - rotary_cos_sin = self.rotary_emb.rotary_cos_sin - supported_rope_cache = ( - rotary_cos_sin.is_cuda - and rotary_cos_sin.dtype == torch.float32 - and rotary_cos_sin.is_contiguous() - and rotary_cos_sin.dim() == 3 - and tuple(rotary_cos_sin.shape[1:]) == (2, 32) - ) - if ( - not supported_main_cache - or not supported_index_cache - or not supported_rope_cache - 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() < num_tokens - ): - return _producer_miss( - covered and layer_uses_nvfp4, - num_tokens, - "the layer's NVFP4 / index caches, RoPE table, KV quant scale or cache slots are not the " - "layout the producer is built for", - ) - - common_args = ( - out_cache_loc[:num_tokens], - self.num_heads, - self.num_key_value_heads, - self.sparse_num_index_heads, - self.head_dim, - rotary_dim, - norm_eps, - self.q_norm.weight, - self.k_norm.weight, - self.index_q_norm.weight, - self.index_k_norm.weight, - rotary_cos_sin, - position_ids.reshape(-1).contiguous().to(torch.int32), - ) - if layer_uses_nvfp4 and producer_pw_supported( - self.num_heads, - self.num_key_value_heads, - self.sparse_num_index_heads, - buffers, - scale_buffers, - ): - q, index_q = m3_producer( - packed.contiguous(), - buffers.view(torch.uint8), - scale_buffers.view(torch.uint8), - index_k_cache, - common_args[0], - kv_quant_scale, - norm_eps, - *common_args[7:12], - common_args[12], - self.num_heads, - get_env_enable_pdl(), - PRODUCER_EARLY_TRIGGER, - ) - elif layer_uses_nvfp4: - if covered: - dispatch.missed( - dispatch.PRODUCER.path, - num_tokens, - f"the NVFP4 pool {tuple(buffers.shape)} / scale pool " - f"{tuple(scale_buffers.shape)} is not the producer's layout", - ) - q, index_q = torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( - packed.contiguous(), - buffers.view(torch.uint8), - scale_buffers, - index_k_cache, - common_args[0], - kv_quant_scale, - *common_args[1:], - ) - else: - q, index_q = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( - packed.contiguous(), - buffers, - index_k_cache, - *common_args, - ) - return q.flatten(1), index_q.flatten(1) - - def _msa_attention_core(self, q, k, v, idx_q, idx_k, attn_metadata, output): - """The built-in MSA attention core, except a pure generation step of a sparse layer whose caches the - fused producer already wrote: that one runs this target's selection and sparse decode.""" - span = attn_metadata.msa_decode_span - if ( - self.is_sparse_attention_layer - and k is None - and attn_metadata.num_contexts == 0 - and span is not None - and span.row_first == 0 - ): - self._decode_sparse_core(q, idx_q, attn_metadata, output) - return output - return super()._msa_attention_core(q, k, v, idx_q, idx_k, attn_metadata, output) - - def _decode_sparse_core(self, q, idx_q, attn_metadata, output) -> None: - """Top-16 block selection and the block-sparse decode of one pure generation step. - - The selection is the identity when the metadata staged it (every request within 16 blocks). Otherwise a span - of at most ``dispatch.FUSED_SCORE_SELECT_MAX_REQUESTS`` requests scores and selects in one launch; a wider - span runs the CuTe DSL scorer (the per-index-head block maxima, reduced to the KV head) and the long-row - selector (the built-in selector beyond its shapes). - """ - from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( - MSA_REQUIRED_TOPK, - msa_paged_kv, - select_blocks_from_maxscore, - ) - - from .kernels.long_select import long_select_blocks - from .kernels.triton_sparse_decode import minimax_m3_sparse_attn_decode - - md = attn_metadata - attn = self.attn - config = attn.m3_config - num_tokens = int(q.shape[0]) - rows = int(md.num_seqs) - qlen = int(md.msa_decode_span.query_len) - block_table = md.msa_block_table[:rows] - seq_lens = md.msa_seq_lens_cuda[:rows] - num_kv_heads = int(attn.indexer.config.num_kv_heads) - # The index scorer's kernels are compiled and loaded by build_targets (prebuild_index_scorer), so no step - # compiles one and the identity can be taken from the first decode step. - identity = md._msa_short_ok and md.msa_short_table is not None - fused = None if identity else self._fused_score_select(idx_q, md, rows, qlen, num_tokens) - if identity: - topk = md.msa_short_table[: num_kv_heads * num_tokens * MSA_REQUIRED_TOPK].view( - num_kv_heads, num_tokens, MSA_REQUIRED_TOPK - ) - elif fused is not None: - topk = fused.permute(1, 0, 2) - else: - idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) - idx_k_cache = md.msa_idx_k_cache(self.layer_idx) - max_score = md.msa_proxy_max_score_view( - config.num_index_heads, md.msa_worst_case_max_k_tiles, num_tokens - ) - early = rows <= INDEX_SCORE_EARLY_MAX_REQUESTS - key = _scorer_key( - idx_q_view.dtype, - config.num_index_heads, - qlen, - idx_q_view.shape[-1], - INDEX_SCORE_SPLIT_K, - early, - ) - if key not in _scorer_runner().kernel_cache: - # Outside require mode the op below compiles it (the only scorer this step has). - dispatch.missed( - INDEX_SCORE_PATH, - num_tokens, - f"the index scorer {key} was not compiled before warmup", - ) - 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, - ) - n_valid_blocks = md.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 - ) - 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=num_tokens <= dispatch.SELECT_EARLY_MAX_ROWS, - ) - if table is None: - dispatch.uncovered( - dispatch.LONG_SELECT.path, - f"{max_score_kv.shape[1]} KV blocks, outside {dispatch.LONG_SELECT_MIN_BLOCKS}.." - f"{dispatch.LONG_SELECT_MAX_BLOCKS}: the built-in selector", - ) - table = select_blocks_from_maxscore( - max_score_kv, - topk=MSA_REQUIRED_TOPK, - n_valid_blocks=n_valid_blocks, - init_blocks=config.init_blocks, - local_blocks=config.local_blocks, - head_major_output=True, - ) - topk = table.permute(1, 0, 2) - manager = md.kv_cache_manager - k_paged, v_paged = msa_paged_kv(manager, self.layer_idx) - nvfp4_args = {} - if manager.is_nvfp4_layer(self.layer_idx): - k_scale, v_scale = _aligned_nvfp4_dequant_scales(attn, self.qkv_proj.kv_scales) - scales = manager.get_block_scale_buffers(self.layer_idx, "HND") - nvfp4_args = dict( - k_block_scale=scales[:, 0], - v_block_scale=scales[:, 1], - k_global_scale=k_scale, - v_global_scale=v_scale, - ) - minimax_m3_sparse_attn_decode( - q.view(num_tokens, self.num_heads, self.head_dim), - k_paged, - v_paged, - topk, - block_table, - seq_lens, - sm_scale=(self.head_dim**-0.5) / float(attn.q_scaling), - output=output.view(num_tokens, self.num_heads, self.head_dim), - decode_query_len=qlen, - identity_blocks=identity, - **nvfp4_args, - ) - - def _fused_score_select(self, idx_q, md, rows: int, qlen: int, num_tokens: int): - """The top-16 table [total_q, num_kv_heads, 16] of a covered decode span from the fused score + select, or None - (after a miss when the span is covered) for the scorer + selector.""" - from .kernels import m3_fused_score_select as fss - - path = dispatch.FUSED_SCORE_SELECT.path - if not dispatch.FUSED_SCORE_SELECT.covers(rows): - dispatch.uncovered( - path, - f"more than {dispatch.FUSED_SCORE_SELECT_MAX_REQUESTS} requests: scorer + selector", - ) - return None - config = self.attn.m3_config - idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) - idx_k_pages = md.msa_idx_k_cache(self.layer_idx).squeeze(1) - max_tiles = int(md.msa_worst_case_max_k_tiles) - n_valid = md.msa_n_valid_blocks - reason = fss.supported(idx_q_view, idx_k_pages, qlen, config.num_kv_heads, max_tiles) - if reason is None and ( - num_tokens != rows * qlen - or n_valid is None - or n_valid.dtype != torch.int32 - or int(n_valid.shape[0]) < num_tokens - ): - reason = "the step's query rows or valid-block counts are not one decode span" - if reason is not None: - dispatch.missed(path, rows, reason) - return None - bufs = fss.buffers(idx_q.device, max_tiles) - if bufs is None: - dispatch.capture_missed(path, rows, "the fused score + select scratch") - return None - return fss.fused_score_select( - idx_q_view, - idx_k_pages, - md.msa_block_table[:rows], - md.msa_seq_lens_cuda[:rows], - n_valid[:num_tokens].contiguous(), - decode_query_len=qlen, - num_kv_heads=config.num_kv_heads, - init_blocks=config.init_blocks, - local_blocks=config.local_blocks, - max_tiles=max_tiles, - scratch=bufs[0], - counters=bufs[1], - ) - - def _dense_decode_fp8_q(self, attn_metadata) -> bool: - """On decode-only steps of a hybrid NVFP4 cache, the dense layers (whose KV pool is FP8) take FP8 q/k/v - straight from the fused norm/RoPE kernel instead of casting them separately.""" - return ( - self.main_kv_is_nvfp4 - and not self.is_sparse_attention_layer - and self._msa_backend_active() - and attn_metadata.num_contexts == 0 - ) - - def _dense_forward( - self, - position_ids: torch.IntTensor, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - all_reduce_params: Optional[AllReduceParams] = None, - **kwargs, - ) -> torch.Tensor: - """The built-in dense attention (layers 0-2), FP8 q/k/v on decode steps. - - Mirrors the built-in ``MiniMaxM3Attention._dense_forward``. - """ - if attn_metadata is None: - raise RuntimeError( - f"MiniMax-M3 dense forward (layer {self.layer_idx}) requires attn_metadata; received None." - ) - qkv = self.qkv_proj(hidden_states) - fused_qkv = self._fused_qk_norm_rope( - qkv, - position_ids, - num_heads_q=self.num_heads, - num_heads_k=self.num_key_value_heads, - num_heads_v=self.num_key_value_heads, - head_dim=self.head_dim, - q_norm=self.q_norm, - k_norm=self.k_norm, - out_fp8=self._emit_fp8_main_qkv() or self._dense_decode_fp8_q(attn_metadata), - ) - if fused_qkv is not None: - q, k, v = self._split_main_qkv(fused_qkv) - else: - assert not self._expect_fused_qk_norm_rope(position_ids), ( - f"MiniMax-M3 dense attention (layer {self.layer_idx}) expected the " - f"fused QK-norm+RoPE kernel (bf16 activations, head_dim=" - f"{self.head_dim}) but fell back to the separate path; qkv dtype " - f"is {qkv.dtype} (expected {self.attn_activation_dtype})." - ) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = self.apply_qk_norm(q, k) - if self.rotary_emb is not None and position_ids is not None: - q, k = self.rotary_emb(position_ids, [q, k]) - o = self._forward_attention_core(q, k, v, None, None, attn_metadata) - return self.o_proj(o, all_reduce_params=all_reduce_params) - - def _sparse_forward( - self, - position_ids: torch.IntTensor, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - all_reduce_params: Optional[AllReduceParams] = None, - **kwargs, - ) -> torch.Tensor: - """The built-in sparse attention, with the fused decode kernels on the steps they take. - - Mirrors the built-in ``MiniMaxM3Attention._sparse_forward``. - """ - if attn_metadata is None: - raise RuntimeError( - f"MiniMax-M3 sparse forward (layer {self.layer_idx}) requires attn_metadata; received None." - ) - packed_qkv = None - packed_idx_qk = None - self._check_producer_projection() - if self.enable_fused_qkv_index_projection: - if ( - self.register_to_config - and is_torch_compiling() - and isinstance(self.attn, MiniMaxM3MsaSparseAttention) - and self._emit_fp8_main_qkv() - and self.attn.indexer_kv_dtype == "fp8" - ): - kv_cache, index_k_cache = attn_metadata.msa_layer_cache_tensors[self.layer_idx] - q, idx_q = torch.ops.trtllm.minimax_m3_fused_sparse_qkv_producer( - hidden_states, - position_ids, - kv_cache, - index_k_cache, - attn_metadata.msa_out_cache_loc, - self.layer_idx_str, - ) - o = self._forward_attention_core(q, None, None, idx_q, None, attn_metadata) - return self.o_proj(o, all_reduce_params=all_reduce_params) - if self.register_to_config and (is_torch_compiling() or is_in_breakable_cuda_graph()): - packed = torch.ops.trtllm.minimax_m3_qkv_index_proj( - hidden_states, position_ids, self.layer_idx_str - ) - num_tokens = ( - position_ids.shape[-1] if position_ids is not None else hidden_states.shape[0] - ) - output = packed.new_empty( - (num_tokens, self.num_heads * self.head_dim), dtype=self.attn_activation_dtype - ) - maybe_bcg_minimax_m3_attn_custom_op_inplace( - None, None, None, None, None, packed, position_ids, self.layer_idx_str, output - ) - return self.o_proj(output, all_reduce_params=all_reduce_params) - packed = self.qkv_proj(hidden_states) - horizontal = self._fused_fp8_qkv_indexer_norm_rope_kv_insert( - packed, position_ids, attn_metadata - ) - if horizontal is not None: - q, idx_q = horizontal - o = self._forward_attention_core(q, None, None, idx_q, None, attn_metadata) - return self.o_proj(o, all_reduce_params=all_reduce_params) - main_size = self.q_size + 2 * self.kv_size - packed_qkv, packed_idx_qk = packed.split( - [main_size, self.index_q_size + self.index_k_size], dim=-1 - ) - - # The compatibility/fallback main and index branches, overlapped on the aux stream. - def _main_norm_rope(): - qkv = packed_qkv if packed_qkv is not None else self.qkv_proj(hidden_states) - fused_qkv = self._fused_qk_norm_rope( - qkv, - position_ids, - num_heads_q=self.num_heads, - num_heads_k=self.num_key_value_heads, - num_heads_v=self.num_key_value_heads, - head_dim=self.head_dim, - q_norm=self.q_norm, - k_norm=self.k_norm, - out_fp8=self._emit_fp8_main_qkv(), - ) - if fused_qkv is not None: - return self._split_main_qkv(fused_qkv) - assert not self._expect_fused_qk_norm_rope(position_ids), ( - f"MiniMax-M3 sparse attention (layer {self.layer_idx}) expected the fused QK-norm+RoPE kernel" - ) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = self.apply_qk_norm(q, k) - if self.rotary_emb is not None and position_ids is not None: - q, k = self.rotary_emb(position_ids, [q, k]) - return q, k, v - - def _index_norm_rope(): - idx_qk = ( - packed_idx_qk if packed_idx_qk is not None else self.index_qk_proj(hidden_states) - ) - graph_fp8_indexer = ( - self.register_to_config - and is_torch_compiling() - and isinstance(self.attn, MiniMaxM3MsaSparseAttention) - and self.attn.indexer_kv_dtype == "fp8" - ) - if not graph_fp8_indexer: - fp8_idx_q = self._fused_fp8_index_qk_norm_rope(idx_qk, position_ids, attn_metadata) - if fp8_idx_q is not None: - return fp8_idx_q, None - fused_idx = self._fused_qk_norm_rope( - idx_qk, - position_ids, - num_heads_q=self.sparse_num_index_heads, - num_heads_k=1, - num_heads_v=0, - head_dim=self.sparse_index_dim, - q_norm=self.index_q_norm, - k_norm=self.index_k_norm, - out_fp8=graph_fp8_indexer, - ) - if fused_idx is not None: - return self._split_index_qk(fused_idx) - assert not self._expect_fused_qk_norm_rope(position_ids), ( - f"MiniMax-M3 sparse index branch (layer {self.layer_idx}) expected the fused QK-norm+RoPE kernel" - ) - idx_q, idx_k = idx_qk.split([self.index_q_size, self.index_k_size], dim=-1) - idx_q, idx_k = self.apply_index_qk_norm(idx_q, idx_k) - if self.rotary_emb is not None and position_ids is not None: - idx_q, idx_k = self.rotary_emb(position_ids, [idx_q, idx_k]) - return idx_q, idx_k - - (q, k, v), (idx_q, idx_k) = maybe_execute_in_parallel( - _main_norm_rope, - _index_norm_rope, - self.ln_events[0], - self.ln_events[1], - self.aux_stream, - disable_on_compile=True, - ) - o = self._forward_attention_core(q, k, v, idx_q, idx_k, attn_metadata) - return self.o_proj(o, all_reduce_params=all_reduce_params) - - -def _plain_capture(spec_metadata, layer_id: int, hidden_states, residual) -> None: - """The in-tree Eagle3 hidden-state capture (no fused fc_norm, no early fc fork).""" - capture_hidden_states(spec_metadata, layer_id, hidden_states, residual, None) - - -class M3DecoderLayer(DecoderLayer): - """One MiniMax-M3 transformer block on this target's modules. - - Mirrors the built-in ``MiniMaxM3DecoderLayer``; the layer boundaries add the - multicast all-reduce + residual + RMSNorm up to its token limits. - """ - - def __init__( - self, - model_config: ModelConfig[PretrainedConfig], - layer_idx: int, - aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - ) -> None: - super().__init__() - config = model_config.pretrained_config - self.hidden_size = config.hidden_size - self.layer_idx = layer_idx - self.mapping = model_config.mapping - self.enable_attention_dp = self.mapping.enable_attention_dp - - _, sparse_layer_ids = get_sparse_layer_ids(config) - disable_index_value_ids = set(get_sparse_disable_index_value_layer_ids(config)) - self.self_attn = M3Attention( - model_config=model_config, - layer_idx=layer_idx, - is_sparse_attention_layer=layer_idx in sparse_layer_ids, - disable_index_value=layer_idx in disable_index_value_ids, - aux_stream=aux_stream_dict[AuxStreamType.Attention], - ) - - _, moe_layer_ids = get_moe_layer_ids(config) - if layer_idx in moe_layer_ids: - self.block_sparse_moe = M3MoE( - model_config=model_config, aux_stream_dict=aux_stream_dict, layer_idx=layer_idx - ) - self.mlp = None - else: - dense_intermediate = int( - getattr(config, "dense_intermediate_size", config.intermediate_size) - ) - self.mlp = _dense_mlp(model_config, dense_intermediate, layer_idx) - self.block_sparse_moe = None - - # Plain (non-Gemma) boundary norms, so they can drive the fused all-reduce - # + residual + RMSNorm epilogue; the loader folds Gemma's (1 + weight) in. - self.input_layernorm = RMSNorm( - hidden_size=config.hidden_size, - eps=config.rms_norm_eps, - dtype=config.torch_dtype, - use_gemma=False, - ) - self.post_attention_layernorm = RMSNorm( - hidden_size=config.hidden_size, - eps=config.rms_norm_eps, - dtype=config.torch_dtype, - use_gemma=False, - ) - - # tp8 without attention DP: every boundary has a cross-rank reduction to fold. - self.pre_feed_forward_fusion = True - self.post_feed_forward_fusion = True - self.allreduce = AllReduce( - mapping=model_config.mapping, - strategy=model_config.allreduce_strategy, - dtype=config.torch_dtype, - ) - # Wired after the weight load to the next layer's input_layernorm (or the final norm). - self.next_layer_layernorm: Optional[RMSNorm] = None - self.next_is_decoder_layer = False - # The Eagle3 hidden-state capture of this layer; the target shell points it at its worker's (fused fc_norm, - # early fc fork). - self.capture: Callable[..., None] = _plain_capture - - def forward( - self, - position_ids: torch.IntTensor, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - residual: Optional[torch.Tensor], - spec_metadata: Optional[SpecMetadata] = None, - lora_params: Optional[dict] = None, - **kwargs, - ) -> Tuple[torch.Tensor, torch.Tensor]: - if residual is None: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - hidden_states = self.self_attn( - position_ids=position_ids, - hidden_states=hidden_states, - attn_metadata=attn_metadata, - all_reduce_params=AllReduceParams(enable_allreduce=False), - lora_params=lora_params, - **kwargs, - ) - capture = spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx) - if self.block_sparse_moe is not None: - return self.forward_moe( - hidden_states, attn_metadata, residual, capture, spec_metadata, lora_params - ) - return self.forward_mlp(hidden_states, residual, capture, spec_metadata, lora_params) - - def _can_mc(self, hidden_states: torch.Tensor, folded: bool = False) -> bool: - """Whether the multicast all-reduce + residual + RMSNorm takes this boundary.""" - return ( - isinstance(hidden_states, torch.Tensor) - and hidden_states.dim() == 2 - and hidden_states.shape[0] - <= (m3_mc_allreduce.MAX_TOKENS if folded else MC_AR_MAX_TOKENS) - and hidden_states.dtype == torch.bfloat16 - and hidden_states.is_contiguous() - ) - - def mc_allreduce_norm(self, x, x2, residual, norm): - """The multicast all-reduce + optional second addend + residual + RMSNorm, or None without a workspace.""" - workspace = m3_mc_allreduce.get_workspace(self.mapping, x.shape[-1]) - if workspace is None: - return None - return m3_mc_allreduce.m3_mc_allreduce_norm( - x, - residual.contiguous(), - norm.weight, - norm.variance_epsilon, - workspace, - use_pdl=get_env_enable_pdl(), - x2=x2, - ) - - def native_allreduce_norm(self, hidden_states, residual, norm): - """The built-in fused all-reduce + residual + RMSNorm.""" - return self.allreduce( - hidden_states, - all_reduce_params=AllReduceParams( - fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, - residual=residual, - norm_weight=norm.weight, - eps=norm.variance_epsilon, - trigger_completion_at_end=False, - ), - ) - - def _pre_feed_forward_norm(self, hidden_states, residual): - if self._can_mc(hidden_states): - fused = self.mc_allreduce_norm( - hidden_states, None, residual, self.post_attention_layernorm - ) - if fused is not None: - return fused - return self.native_allreduce_norm(hidden_states, residual, self.post_attention_layernorm) - - def _next_layer_norm(self, hidden_states, residual, fused_boundary: bool): - """The next layer's input norm at the boundary; ``fused_boundary`` False keeps the module's own reduction.""" - if self.next_layer_layernorm is None: - return hidden_states, residual - parts = hidden_states if isinstance(hidden_states, tuple) else (hidden_states, None) - if not fused_boundary: - return self.next_layer_layernorm(parts[0], residual) - if self._can_mc(parts[0], folded=parts[1] is not None): - fused = self.mc_allreduce_norm(parts[0], parts[1], residual, self.next_layer_layernorm) - if fused is not None: - return fused - hidden_states = parts[0] if parts[1] is None else parts[0].add_(parts[1]) - return self.native_allreduce_norm(hidden_states, residual, self.next_layer_layernorm) - - def _fused_feed_forward_boundary(self, capture: bool) -> bool: - """Whether the feed-forward output goes unreduced to the fused boundary (all-reduce + residual + next norm in - one kernel). An Eagle3 capture layer takes it too and captures the boundary's reduced residual sum; only the - last layer (no next norm, its raw output captured) keeps the module's own all-reduce.""" - return not capture or self.next_layer_layernorm is not None - - def _feed_forward_boundary( - self, hidden_states, residual, fused_boundary, capture, spec_metadata - ): - """The layer's exit: the next norm at the boundary, with this layer's Eagle3 capture.""" - if not capture: - return self._next_layer_norm(hidden_states, residual, fused_boundary) - if not fused_boundary: - self.capture(spec_metadata, self.layer_idx, hidden_states, residual) - return self._next_layer_norm(hidden_states, residual, False) - normed, updated = self._next_layer_norm(hidden_states, residual, True) - self.capture(spec_metadata, self.layer_idx, updated, None) - return normed, updated - - def forward_moe( - self, hidden_states, attn_metadata, residual, capture, spec_metadata, lora_params - ): - hidden_states, residual = self._pre_feed_forward_norm(hidden_states, residual) - fused_boundary = self._fused_feed_forward_boundary(capture) - hidden_states = self.block_sparse_moe( - hidden_states, - attn_metadata, - final_all_reduce_params=AllReduceParams(enable_allreduce=False) - if fused_boundary - else None, - lora_params=lora_params, - return_parts=fused_boundary and self._can_mc(hidden_states, folded=True), - ) - return self._feed_forward_boundary( - hidden_states, residual, fused_boundary, capture, spec_metadata - ) - - def forward_mlp(self, hidden_states, residual, capture, spec_metadata, lora_params): - hidden_states, residual = self._pre_feed_forward_norm(hidden_states, residual) - fused_boundary = self._fused_feed_forward_boundary(capture) - hidden_states = self.mlp( - hidden_states, - final_all_reduce_params=AllReduceParams(enable_allreduce=False) - if fused_boundary - else None, - lora_params=lora_params, - ) - return self._feed_forward_boundary( - hidden_states, residual, fused_boundary, capture, spec_metadata - ) - - -class M3TextModel(DecoderModel): - """The MiniMax-M3 text decoder on this target's layers. - - Its forward belongs to the core in ``modeling.py``, which dispatches each - step to the prefill or the decode target. - """ - - def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: - _validate_sparse_attention_runtime_config(model_config) - super().__init__(model_config) - config = model_config.pretrained_config - self.vocab_size = config.vocab_size - self.aux_stream_dict = { - AuxStreamType.Attention: torch.cuda.Stream(), - AuxStreamType.MoeShared: torch.cuda.Stream(), - AuxStreamType.MoeChunkingOverlap: torch.cuda.Stream(), - } - self.embed_tokens = Embedding( - config.vocab_size, - config.hidden_size, - dtype=config.torch_dtype, - enable_torch_compile_for_embedding=model_config.enable_torch_compile_for_embedding, - ) - self.layers = nn.ModuleList( - [ - M3DecoderLayer(model_config, layer_idx, self.aux_stream_dict) - for layer_idx in range(config.num_hidden_layers) - ] - ) - # Lets the executor enable the automatic FlashInfer MXFP8 path for the decode graphs. - self._use_flashinfer_mxfp8_decode_graph_default = True - self.norm = RMSNorm( - hidden_size=config.hidden_size, - eps=config.rms_norm_eps, - dtype=config.torch_dtype, - use_gemma=False, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/linears.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/linears.py deleted file mode 100644 index 07c0e668c866..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/linears.py +++ /dev/null @@ -1,232 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The decode-graph MXFP8 GEMM dispatch of this target's projections. - -Every MXFP8 projection of the text model (QKV, o_proj, the dense MLPs and the -shared experts) runs the built-in ``MXFP8LinearMethod`` outside decode CUDA -graphs. Inside them (``flashinfer_mxfp8_decode_graph_capture_active``) the -projection's (N, K) and the step's rows pick one path (``dispatch.linear_path``, -row ranges per projection in ``dispatch.py``): - -1. an activation its producer already quantized to MXFP8 (the boundary - all-reduce or the attention merge publish it in ``m3_prequant.SLOT``): the - tcgen05 decode GEMV on the rows it serves, else the decode-graph MXFP8 GEMM - (AutoTuner-selected backend per token bucket, the built-in method's GEMM); -2. a bf16 activation on the rows the GEMV serves quantizing in its prologue; -3. otherwise the built-in method (MXFP8 quantize + GEMM, AutoTuner-selected per bucket). - -No path runs FlashInfer's own per-rank config sweep: its tactics are not merged -across ranks and it profiled wherever the projection ran (the MoE aux stream for -the shared expert), one CUDA graph per tactic. - -A kernel declining rows its path covers is a miss (``dispatch.missed``); the -chain then continues down the list. The GEMVs push their rows straight into -the boundary all-reduce the decode forward armed for the projection's push -slot (``m3_push``, fixed at construction). -""" - -from __future__ import annotations - -from typing import List, Optional, Tuple - -import torch -from torch import nn - -from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_gemm_autotuned -from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl -from tensorrt_llm._torch.modules.linear import ( - Linear, - MXFP8LinearMethod, - flashinfer_mxfp8_autotune_active, - flashinfer_mxfp8_decode_graph_capture_active, -) -from tensorrt_llm._torch.utils import is_torch_compiling - -from . import dispatch -from .kernels import m3_mc_allreduce, m3_prequant -from .kernels.mx_gemv import op as mx_gemv -from .ops.gemm.m3_mx_gemv import m3_mx_gemv -from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 - -# The boundary all-reduce a projection's GEMV pushes its rows into, when the decode forward arms one: the attention -# boundary (o_proj) or the MoE boundary's shared-expert set (the shared expert's down projection). -PUSH_ATTN = "attn" -PUSH_MOE = "moe" - - -def mark_m3_linear(module: Linear, push: Optional[str] = None) -> None: - """Mark a projection for this target's decode dispatch; ``adopt_m3_linears`` applies it after the load. - - ``push``: ``PUSH_ATTN`` / ``PUSH_MOE`` when its GEMV pushes into that boundary's all-reduce whenever the decode - forward armed it. - """ - module.m3_push = push - - -def adopt_m3_linears(root: nn.Module) -> None: - """Put every marked projection under ``root`` on the decode dispatch. - - Runs once the weights are real: the quantization method a projection ends - up with is only final after the model applied its quantization config. - It also compiles every decode GEMV kernel a projection can reach here, - before the executor's warmup: on the current (main) stream, the same shapes - in the same order on every rank, and never inside a forward between - collectives or on the MoE aux stream (``prebuild_m3_gemvs``). - """ - use_pdl = get_env_enable_pdl() - compiled = False - for name, module in root.named_modules(): - if not hasattr(module, "m3_push"): - continue - method = module.quant_method - if not isinstance(method, MXFP8LinearMethod) or not method.use_cutlass: - raise ValueError( - f"{name}: the MiniMax-M3 decode dispatch needs a CUTLASS MXFP8 projection, got {type(method).__name__}" - ) - if type(method) is MXFP8LinearMethod: - method.__class__ = M3MXFP8LinearMethod - mx_gemv.ensure_compiled(module.weight, module.weight_scale, use_pdl) - prebuild_m3_gemvs(module, use_pdl) - compiled = True - if compiled: - torch.cuda.synchronize() - - -def _bucket_tokens(rows: int) -> int: - """The token capacity of the multicast all-reduce bucket a ``rows``-token boundary uses.""" - return next(b for b in m3_mc_allreduce.BUCKETS if rows <= b) - - -def gemv_variants(module: Linear) -> List[Tuple[int, bool, Optional[Tuple[int, int, int]]]]: - """Every (rows, bf16 activation, push (world, bucket tokens, mode)) call the decode dispatch can route to this - projection's GEMV (``M3MXFP8LinearMethod.apply``): 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.""" - n, k = module.weight.shape - push_rows = {PUSH_ATTN: dispatch.ATTN_PUSH.keys, PUSH_MOE: dispatch.MOE_PUSH.keys}.get( - module.m3_push, frozenset() - ) - world = int(module.tp_size) - variants = [] - for rows in range(1, dispatch.DECODE_MAX_TOKENS + 1): - for x_bf16, served in ( - (False, rows in dispatch.MX_GEMV_ROWS.get((n, k), ())), - (True, rows <= dispatch.MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0)), - ): - if not served: - continue - variants.append((rows, x_bf16, None)) - if rows in push_rows: - variants.append( - (rows, x_bf16, (world, _bucket_tokens(rows), dispatch.AR_PUSH_MODE)) - ) - return variants - - -def prebuild_m3_gemvs(module: Linear, use_pdl: bool) -> int: - """Compile every GEMV kernel of ``gemv_variants`` (the push ones included) without launching; returns how many - distinct kernels the projection takes.""" - keys = set() - for rows, x_bf16, push in gemv_variants(module): - key = mx_gemv.precompile( - module.weight, module.weight_scale, rows, x_bf16, use_pdl=use_pdl, ar=push - ) - if key is not None: - keys.add(key) - return len(keys) - - -def _push_target(module: Linear): - """The armed boundary all-reduce this projection pushes its output into, if any.""" - if module.m3_push == PUSH_MOE: - return m3_prequant.MOE_PUSH.take_shared() - if module.m3_push == PUSH_ATTN: - return m3_prequant.AR_PUSH.take_target() - return None - - -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 _note_pushed(output: Optional[torch.Tensor], module: Linear) -> None: - if output is None or not mx_gemv.last_pushed: - return - if module.m3_push == PUSH_MOE: - m3_prequant.MOE_PUSH.shared_pushed = output.data_ptr() - else: - m3_prequant.AR_PUSH.pushed = output.data_ptr() - - -class M3MXFP8LinearMethod(MXFP8LinearMethod): - """``MXFP8LinearMethod`` with the decode-graph dispatch in the module docstring.""" - - def apply(self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor]): - if bias is not None or input.dim() != 2: - return super().apply(module, input, bias) - input = input.contiguous() - weight, weight_scale = module.weight, module.weight_scale - use_pdl = get_env_enable_pdl() - decode_graph = not is_torch_compiling() and flashinfer_mxfp8_decode_graph_capture_active() - if not decode_graph: - return super().apply(module, input, bias) - - rows = input.shape[0] - n, k = weight.shape - prequant = m3_prequant.SLOT.take(input) - if module.dtype != torch.bfloat16 or (prequant is None and input.dtype != torch.bfloat16): - path = ( - dispatch.PREQUANT_GEMM_PATH if prequant is not None else dispatch.GENERIC_GEMM_PATH - ) - else: - path = dispatch.linear_path(n, k, rows, prequant is not None) - output = None - if prequant is not None: - act, act_scale = prequant - if path == dispatch.MX_GEMV_PATH: - output = m3_mx_gemv( - act, act_scale, weight, weight_scale, use_pdl=use_pdl, ar=_push_target(module) - ) - _note_pushed(output, module) - if output is None: - dispatch.missed( - path, - rows, - f"the GEMV declined the [{n}, {k}] projection's MXFP8 activation" - + _uncompiled_note(), - ) - if output is None: - if path == dispatch.PREQUANT_GEMM_PATH: - dispatch.uncovered( - dispatch.MX_GEMV_PATH, - f"the [{n}, {k}] projection at rows the GEMV does not serve", - ) - output = mxfp8_gemm_autotuned( - act, - act_scale, - weight, - weight_scale, - module.dtype, - tune=flashinfer_mxfp8_autotune_active(), - ) - return output - if path == dispatch.MX_GEMV_BF16_PATH: - output = m3_mx_gemv_bf16( - input, weight, weight_scale, use_pdl=use_pdl, ar=_push_target(module) - ) - _note_pushed(output, module) - if output is None: - dispatch.missed( - path, - rows, - f"the GEMV declined the [{n}, {k}] projection's input" + _uncompiled_note(), - ) - if output is None: - if path == dispatch.GENERIC_GEMM_PATH: - dispatch.uncovered( - dispatch.MX_GEMV_BF16_PATH, - f"the [{n}, {k}] projection's bf16 input beyond the decode GEMV's rows", - ) - # Stock's decode-graph dispatch for this shape (mxfp8_quantize_gemm_autotuned): same tactics, same bits. - output = super().apply(module, input, bias) - return output diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py index de2214dfddfa..e2eedd3a8a17 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py @@ -1,559 +1,4204 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""ModelingV2 target: MiniMax-M3 NVFP4 / sm_103 / tp8. - -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) -plus one shared expert. The checkpoint is a modelopt mixed-precision export: -NVFP4 routed experts, MXFP8 projections, bf16 router, embedding and head; the KV -cache is NVFP4 for the sparse layers and FP8 for the dense ones. - -``tp8`` is eight GB300s over two trays (multicast all-reduces over MNNVL), -head-split attention (8 query heads and 1 KV head, duplicated, per rank) and -tensor-parallel experts (every expert on every rank, width 384 per rank), no -attention DP. Speculative decoding, when configured, is one-model Eagle3 with -draft length 3: the in-tree BF16 drafter over the full vocabulary and the BF16 -target head, both at the checkpoint's precision; decode steps of up to 8 -requests verify and draft through fused kernels at that precision -(``spec_worker.py``). The decode kernels are specialized for one request -verifying a draft (4 query rows), the deployment this target was measured on. - -**Each step is dispatched once, on the host, from its shape** (``_select_target``): - -* ``DecodeTarget`` takes a pure generation step of at most 48 tokens (up to 8 - requests, each one token or a draft verify of 1 + 3 + extra rows). It runs - the layers with the decode-step kernels of ``layers.py`` and chooses each - layer boundary's all-reduce by the step's token bucket (``DECODE_PLANS``). - Every fast path it can take, the rows each covers and what a step that cannot - take one does (raise under ``modeling_v2='require'``, else log once and fall - back) are in ``dispatch.py``. -* ``PrefillTarget`` takes every other step (prefill, mixed steps, wider decode - steps) on the built-in modules' generic path. - -**What this target asserts rather than adapts**: SM 10.3, the topology above, -the checkpoint shape and quantization (``routing.py``), the MiniMax-M3 MSA -sparse attention backend with its NVFP4 hybrid KV cache, and text-only input: -the vision tower is not built and a step carrying image input raises. +"""ModelingV2 target: MiniMax-M3 NVFP4 / sm_103 / 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 GB300s over two trays (multicast all-reduces over MNNVL): head-split attention (8 query heads, one +KV head and one index head per rank; the 4 KV and index heads each duplicated on 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 8 requests verify and draft through fused kernels at that precision +(``M3Eagle3Worker``). The decode kernels are specialized for one request verifying a draft (4 query rows), the +deployment this target was measured on. + +**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 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 for the +o_proj. It publishes it in ``kernels/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 typing import Any, Dict, List, Optional +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.distributed import AllReduceParams +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.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_minimaxm3 import ( - MiniMaxM3ForCausalLM, - get_text_model_config, - is_minimax_m3_vl_config, +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.models.modeling_speculative import SpecDecOneEngineForCausalLM -from tensorrt_llm._torch.modules.linear import 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.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.functional import PositionEmbeddingType from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantAlgo -from . import dispatch from . import weights as _weights -from .dispatch import AR_PUSH_MODE, EMBED_FIRSTNORM_MAX_TOKENS, DecodePlan, plan_for -from .kernels import m3_mc_allreduce, m3_prequant +from .kernels import m3_draft_step as m3_draft_step_kernel +from .kernels import m3_fused_score_select as _fused_score_select_kernel +from .kernels import m3_index_score as _index_score_kernel +from .kernels import m3_mc_allreduce, m3_prequant, m3_spec_tail +from .kernels import m3_sampled_accept as m3_sampled_accept_kernel +from .kernels import m3_shard_accept as m3_shard_accept_kernel +from .kernels import msa_kv_update as _msa_kv_update_kernel +from .kernels.long_select import long_select_blocks +from .kernels.m3_router_gemv import m3_router_gemv as _router_gemv +from .kernels.m3_router_gemv import m3_router_gemv_fp4 as _router_gemv_fp4 +from .kernels.m3_router_gemv import m3_router_tc as _router_tc from .kernels.mx_gemv import op as mx_gemv -from .layers import M3DecoderLayer, M3MoE, M3TextModel, prebuild_index_scorer -from .linears import adopt_m3_linears +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_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_fused_score_select import ( + MAX_REQUESTS as FUSED_SCORE_SELECT_KERNEL_MAX_REQUESTS, +) +from .ops.attention.m3_fused_score_select import m3_fused_score_select +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_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 MAX_BLOCKS as SELECT_LONG_MAX_BLOCKS +from .ops.attention.m3_select_long import m3_select_long +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.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_draft_lowm import TACTICS as DRAFT_LOWM_TACTICS +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, 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.moe.m3_moe_engine import M3MoeState, ar_push_desc, m3_moe_engine +from .ops.moe.m3_router_gemv import m3_router_gemv +from .ops.moe.m3_router_gemv_fp4 import m3_router_gemv_fp4 +from .ops.moe.m3_router_tc import m3_router_tc from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm -from .runtime import make_runtime_hooks -from .spec_worker import M3Eagle3Worker, M3SpecDecForwardMixin +from .ops.spec.m3_draft_input import m3_draft_input +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_shard_accept import m3_shard_accept _SM = (10, 3) +# -------------------------------------------------------------------------------------------------------------------- +# Row ranges +# -------------------------------------------------------------------------------------------------------------------- +# The largest generation step the decode kernels take: 8 requests verifying 1 + 3 drafts and their padding rows. A +# wider pure generation step stays on the decode target and runs its generic layers. +DECODE_MAX_TOKENS = 48 # capacity: fixed: the decode kernels' step bound -class MiniMaxM3Tp8ModelingV2(M3TextModel): - """The text model of this target: owns the layers and dispatches each step to a target.""" - def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: - super().__init__(model_config) - self._embed_owners: Dict[tuple, DecodeEmbedding] = {} - self._embed_identity: Optional[tuple] = None - self._prefill: Optional[PrefillTarget] = None - self._decode: Optional[DecodeTarget] = None +@dataclass(frozen=True) +class DecodePlan: + """The layer-boundary collectives of one decode token bucket (GB300 measurements in the comments).""" - def build_targets(self) -> None: - """Build the two targets, once the weights are real.""" - table = self.embed_tokens.weight - norm = self.layers[0].input_layernorm - identity = ( - id(table), - table.data_ptr(), - id(norm.weight), - norm.weight.data_ptr(), - float(norm.variance_epsilon), + 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 (kernels/m3_mc_allreduce.BUCKETS). +DECODE_PLANS = ( + DecodePlan("tiny", 8, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), + DecodePlan("small", 16, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), + DecodePlan("verify", DECODE_MAX_TOKENS, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), +) + +# 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. Two requests are neutral (the grid is unchanged) and keep the step's bound. +DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({1: 16384}) +# The same cap on the Eagle3 drafter layer's trtllm-gen decode (8 query heads over one KV head per rank): GB300, one +# request, 63.7K cached tokens: 22.5 -> 11.7 us at one query row, 24.3 -> 12.5 us at four; 188K tokens: 28.2 -> 20.2 us. +# Two requests are flat and keep the step's bound. +DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({1: 16384}) + +# The sparse layers' NVFP4 QKV + indexer producer (ops/attention/m3_producer, every step its geometry fits) triggers its +# dependent kernel on entry, so the kernel after it launches and stages its pre-wait loads while the producer runs. +PRODUCER_EARLY_TRIGGER = 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 = 16 # capacity: fixed: measured push crossover +AR_PUSH_MODE = 2 +# The 1..8-token MoE engine (router, routed experts at the TP8 width 384 and the shared expert in one launch) pushes +# its sum into the MoE boundary's all-reduce (prepushed mode 2). +MOE_ENGINE_MAX_TOKENS = 8 # capacity: fixed: the engine's built rows +MOE_PUSH_MAX_TOKENS = MOE_ENGINE_MAX_TOKENS # capacity: fixed: the engine's 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. +MC_FINALIZE_MIN_TOKENS = MOE_ENGINE_MAX_TOKENS + 1 # capacity: fixed: the rows above the engine +# The router: the fp32 GEMV at one row, the tensor-core router (fused with the routed experts' NVFP4 activation +# quantization) at 2..48 rows. +ROUTER_GEMV_MAX_TOKENS = 1 # capacity: fixed: measured router crossover +ROUTER_TC_MAX_TOKENS = 48 # capacity: fixed: the tensor-core router's built rows +# 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 (kernels/m3_shard_accept), at the configured draft length. The lm_head +# is the in-tree BF16 GEMM: at the TP8 shard (25008 rows) a coded copy of the head saves no time. +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) +# 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 (M3Eagle3Worker). At TP8 it forks beside the verify head instead: forked early +# its GEMM lands beside the target tail it was meant to hide under. +FC_EARLY_MAX_ROWS = 0 # capacity: fixed: measured fork crossover (never at TP8) +# 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 (kernels/m3_spec_tail.draft_input), the drafter layer's attention, MLP and +# all-reduces, and one kernel for the BF16 draft head, the cross-rank argmax and the pass's bookkeeping +# (kernels/m3_draft_step), at the configured draft length. The head kernel takes one row per request, at most 8, and +# streams the head's BF16 weight fragments. +DRAFT_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN # capacity: fixed: the built draft length +DRAFT_STEP_MAX_REQUESTS = DRAFT_STEP_KERNEL_MAX_TOKENS +# The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows +# (kernels/triton_sparse_decode.MQ_MIN_TOKENS). +SPARSE_DECODE_MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query crossover +# Decode spans of at most this many requests score and select their top-16 blocks in one launch +# (ops/attention/m3_fused_score_select), which loads the older index-K pages before its grid wait; larger spans run +# the index scorer and the long-row selector. +FUSED_SCORE_SELECT_MAX_REQUESTS = FUSED_SCORE_SELECT_KERNEL_MAX_REQUESTS +# 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 = SELECT_LONG_MAX_BLOCKS +LONG_SELECT_KEY_WIDTHS = (2, 4, 9) +# The long-row select releases the sparse decode after it right after its own grid wait up to this many query rows; +# wider steps keep the trigger at the select's exit (the early-resident attention CTAs crowd the select grid). +SELECT_EARLY_MAX_ROWS = 8 # capacity: fixed: measured early-release crossover +# The 8-head long-context sparse decode (ops/attention/m3_sparse_decode_lc8) stages the request length, its query rows +# and the request's block-table row before its grid wait up to this many query rows; above, the staged row's shared +# memory costs occupancy and the kernel stages only after the wait (same results bit for bit). +LC8_PREWAIT_MAX_ROWS = 8 # capacity: fixed: measured occupancy crossover +# The MSA index scorer (ops/attention/m3_index_score) prefetches the older index-K pages into L2 before its grid wait +# at every batch it runs (more than FUSED_SCORE_SELECT_MAX_REQUESTS requests); it never runs its one-request +# early-issue schedule here. +INDEX_SCORE_EARLY_MAX_REQUESTS = 0 # capacity: fixed: the early-issue schedule is off at TP8 + +# 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). +_ALL_ROWS = frozenset(range(1, DECODE_MAX_TOKENS + 1)) +MX_GEMV_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { + QKV_SPARSE: _ALL_ROWS, + QKV: _ALL_ROWS, + O_PROJ: _ALL_ROWS, + DENSE_GATE_UP: _ALL_ROWS, + DENSE_DOWN: _ALL_ROWS, +} +# 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: 48, + DENSE_GATE_UP: 16, + DENSE_DOWN: 32, +} +# The Eagle3 drafter's BF16 linears (qkv, o_proj, gate_up, down, fc) through FlashInfer's fixed-tactic low-M kernels +# (ops/gemm/m3_draft_lowm): per (N, K), the rows its tactic table covers; more rows keep the in-tree GEMM. +DRAFT_LOWM_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { + nk: frozenset(range(1, table[-1][0] + 1)) for nk, table in DRAFT_LOWM_TACTICS.items() +} + +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, +) +# The long-context sparse decode: "prewait" up to LC8_PREWAIT_MAX_ROWS query rows, "plain" above. +SPARSE_DECODE_LC8 = RowRange( + "m3_sparse_decode_lc8", + frozenset({"prewait", "plain"}), + 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. +PRODUCER = RowRange("m3_producer", frozenset({"nvfp4_tp8"}), m3_producer, _const("nvfp4_tp8")) +# Requests of a pure decode span whose blocks are scored and selected in one launch. +FUSED_SCORE_SELECT = RowRange( + "m3_fused_score_select", + _rows(1, FUSED_SCORE_SELECT_MAX_REQUESTS), + m3_fused_score_select, + _spec_rows("batch"), +) +INDEX_SCORE = RowRange( + "m3_index_score", + frozenset({"prefetch"}), + m3_index_score, + lambda spec: () if spec["early"] else ("prefetch",), +) +LONG_SELECT = RowRange( + "long_select", frozenset(LONG_SELECT_KEY_WIDTHS), m3_select_long, _long_select_widths +) +MC_ALLREDUCE = RowRange( + "mc_allreduce_norm", + _rows(1, DECODE_MAX_TOKENS), + 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), +) +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, DECODE_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") +) +ROUTER_GEMV = RowRange( + "m3_router_gemv", _rows(1, ROUTER_GEMV_MAX_TOKENS), m3_router_gemv, _spec_rows("m") +) +ROUTER_GEMV_FP4 = RowRange( + "m3_router_gemv_fp4", _rows(1, ROUTER_GEMV_MAX_TOKENS), m3_router_gemv_fp4, _spec_rows("m") +) +ROUTER_TC = RowRange( + "m3_router_tc", + _rows(ROUTER_GEMV_MAX_TOKENS + 1, ROUTER_TC_MAX_TOKENS), + m3_router_tc, + _spec_rows("m"), +) +# 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_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 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 (), +) + +# (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, + DENSE_DECODE_SPLIT, + SPARSE_DECODE_PRODUCER, + SPARSE_DECODE_LC8, + PRODUCER, + FUSED_SCORE_SELECT, + INDEX_SCORE, + LONG_SELECT, + MC_ALLREDUCE, + ATTN_PUSH, + MOE_PUSH, + MC_FINALIZE, + MOE_ENGINE, + ROUTER_GEMV, + ROUTER_GEMV_FP4, + ROUTER_TC, + MX_GEMV, + MX_GEMV_BF16, + DRAFT_LOWM, + SHARD_ACCEPT, + DRAFT_STEP, + 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: 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 + + +def router_path(num_tokens: int) -> Optional[str]: + """The router kernel when the MoE runs unfused: the fp32 GEMV, the tensor-core router, or None (built-in gate).""" + if ROUTER_GEMV.covers(num_tokens): + return ROUTER_GEMV.path + if ROUTER_TC.covers(num_tokens): + return ROUTER_TC.path + return None + + +MX_GEMV_PATH = MX_GEMV.path +MX_GEMV_BF16_PATH = MX_GEMV_BF16.path +# 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) -> str: + """The decode-graph MXFP8 GEMM of an [N, K] projection on ``num_tokens`` rows. + + ``prequant``: the producer published the activation in MXFP8. + """ + 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}" ) - if self._embed_owners and identity != self._embed_identity: - raise RuntimeError( - "the embedding or the first norm moved after the decode embedding owners were built" - ) - if not self._embed_owners: - self._embed_owners = { - (rows, dtype): DecodeEmbedding(table, norm, rows, dtype) - for rows in range(1, EMBED_FIRSTNORM_MAX_TOKENS + 1) - for dtype in (torch.int32, torch.int64) - } - self._embed_identity = identity - self.prebuild_kernels() - self._prefill = PrefillTarget(self) - self._decode = DecodeTarget(self) + _log_once( + logger.warning, + path, + reason, + f"[M3 dispatch] {path}{at}: {reason}; falling back (logged once per reason)", + ) - 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, 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 ``dispatch.missed``). - Returns the seconds each part took and the scorer keys.""" - times: Dict[str, object] = {} - mx_gemv.freeze(False) - start = time.perf_counter() - adopt_m3_linears(self) - times["gemv_s"] = time.perf_counter() - start - start = time.perf_counter() - for layer in self.layers: - moe = layer.block_sparse_moe - if isinstance(moe, M3MoE): - moe.prebuild_engines() - torch.cuda.synchronize() - times["moe_s"] = time.perf_counter() - start - start = time.perf_counter() - times["scorer_keys"] = prebuild_index_scorer( - self.layers, self.decode_query_lens(), self.embed_tokens.weight.device +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 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() + + +# -------------------------------------------------------------------------------------------------------------------- +# 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 (m3_mc_allreduce.MAX_TOKENS). +MC_AR_MAX_TOKENS = 8 # capacity: fixed: measured crossover vs the one-shot all-reduce +# 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) + +# -------------------------------------------------------------------------------------------------------------------- +# 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 not in the checkpoint): routing declines every other one + (``routing._runtime_mismatch``), so this is the backstop for a construction that did not go through routing. + + 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'" ) - torch.cuda.synchronize() - times["scorer_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)" - % (times["gemv_s"], times["moe_s"], times["scorer_s"], len(times["scorer_keys"])) + if getattr(sparse, "indexer_kv_dtype", None) != "fp8": + raise ValueError( + "this MiniMax-M3 target requires sparse_attention_config.indexer_kv_dtype='fp8'" ) - return times + 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 _select_target(self, attn_metadata: AttentionMetadata, input_ids: Optional[torch.Tensor]): - """Which target runs this step: host-side counts only, fixed per CUDA-graph capture.""" - if ( - attn_metadata.num_contexts == 0 - and input_ids is not None - and input_ids.ndim == 1 - and dispatch.is_decode_step(input_ids.numel()) + +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, "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)), ): - return self._decode - return self._prefill + 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 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, - **kwargs, - ) -> torch.Tensor: - 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") - target = ( - self._select_target(attn_metadata, input_ids) - if inputs_embeds is None - else self._prefill + +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, 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. + """ + 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) + _record("linear", (n, k), rows, path, push is not None) + output = 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: + 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") + 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: + 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() ) - return target.forward( - attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, **kwargs + if path == GENERIC_GEMM_PATH: + uncovered( + MX_GEMV_BF16_PATH, + f"the [{n}, {k}] projection's bf16 input beyond the decode GEMV's rows", ) + return _mxfp8_linear(core, x, w, s, mode), False -class PrefillTarget: - """The general case: context rows, mixed steps and wide generation steps, on the generic layer path.""" +# -------------------------------------------------------------------------------------------------------------------- +# 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.""" + return ( + isinstance(x, torch.Tensor) + and x.dim() == 2 + and x.shape[0] <= (m3_mc_allreduce.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): + """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.""" + 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, + ) + + +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``: the fp32 GEMV at one row, the tensor-core router up to 48, else the + fp32 linear.""" + rows = h.shape[0] if h.dim() == 2 else 0 + path = router_path(rows) + w = core.w[f"l{i}_router"] + if path == ROUTER_GEMV.path: + logits = _router_gemv(h, w, use_pdl=get_env_enable_pdl()) + if logits is not None: + return logits + missed(path, rows, "the router GEMV declined the operands") + elif path == ROUTER_TC.path: + fused = _router_tc(h, w, use_pdl=get_env_enable_pdl()) + if fused is not None: + return fused[0] + missed(path, rows, "the tensor-core router declined the operands") + else: + uncovered(ROUTER_TC.path, f"more than {ROUTER_TC_MAX_TOKENS} rows") + return m3_fp32_linear(h, w) + + +def _router_fp4(core, i: int, h: torch.Tensor): + """Router logits and the routed experts' NVFP4 activation from one kernel, or None (then ``_router_logits``).""" + if h.dim() != 2 or h.shape[0] > ROUTER_TC_MAX_TOKENS: + return None + rows = h.shape[0] + path = router_path(rows) + if path == ROUTER_GEMV.path: + path = ROUTER_GEMV_FP4.path + scale = core.moe[i].fc31_input_scale + if path == ROUTER_TC.path: + fused = _router_tc(h, core.w[f"l{i}_router"], scale, use_pdl=get_env_enable_pdl()) + else: + fused = _router_gemv_fp4(h, core.w[f"l{i}_router"], scale, use_pdl=get_env_enable_pdl()) + if fused is None: + missed(path, rows, "the router declined the operands") + return fused + + +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 + + +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 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.""" + 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) + ): + _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, + ) + 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 _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] = {} + + +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 + 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] = {} + 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", None)) + 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, 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). Returns the + seconds each part took and the scorer keys.""" + times: Dict[str, object] = {} + if not self.mxfp8.quant_method.use_cutlass: + raise ValueError("the MiniMax-M3 decode dispatch needs the CUTLASS MXFP8 GEMM") + 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(), self.w["embed"].device + ) + torch.cuda.synchronize() + times["scorer_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)" + % (times["gemv_s"], times["moe_s"], times["scorer_s"], len(times["scorer_keys"])) + ) + 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: + 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 + + # ------------------------------------------------------------------------------------------ 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 (up to ``INDEX_SCORE_EARLY_MAX_REQUESTS`` requests, + none here) 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) if INDEX_SCORE_EARLY_MAX_REQUESTS else (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 + + # --------------------------------------------------------------------------------------------- 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, + **kwargs, + ) -> torch.Tensor: + 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 + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# 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 + ): + 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) -> torch.Tensor: + """A gated MXFP8 MLP (the dense MLP, 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") + 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") + 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 fused router with the NVFP4 activation up to 48 rows, else the + router logits and the experts quantizing their bf16 input (chunked).""" + core = self.core + fused = _router_fp4(core, i, h) + if fused is not None: + router_logits, x_fp4, x_sf = fused + return _routed_experts(core, i, x_fp4, x_sf, router_logits, True) + router_logits = _router_logits(core, i, h) + return self._routed_chunks(i, h, router_logits) + + 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"), + 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") + 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") + 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") + 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") + 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") + + 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") + 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") + + 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-row engine, pushing its sum into the boundary all-reduce; the unfinalized MoE with the + finalize in the all-reduce from 9 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: MiniMaxM3Tp8ModelingV2) -> None: + self.core = core + self._moe_push_descs: Dict[tuple, Any] = {} + # 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, + metadata=md, + ) + + def forward( + self, attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, capture + ): + 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): + if ( + not kernels + or residual is None + or (spec_metadata is not None and spec_metadata.is_layer_capture(i)) + ): + 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) + 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): + """All-reduce + (second addend) + residual + RMSNorm: the multicast kernel, else the built-in all-reduce.""" + core = self.core + if use_mc: + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w, quant=quant) + if fused is not None: + 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): + fused = _mc_allreduce_norm( + self.core, x, x2, residual, norm_w, quant=quant, prepushed=prepushed + ) + if fused is None: + raise RuntimeError(f"layer {i}: rows were pushed into an all-reduce that did not run") + return fused + + def _layer(self, i: int, plan: DecodePlan, position_ids, hidden, residual, step): + 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: + 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: + out, _ = self._swiglu_mlp(hidden, f"l{i}_mlp") + return self._boundary(plan.attn_mc_ar, out, residual, norm_w) + + 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: + moe_push = self._moe_push_targets(num_tokens) + if 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" + ) + 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: + return self._prepushed(i, shared, residual, norm_w, quant, 2) + else: + if engine_pushed and shared_pushed: + return self._prepushed(i, shared, residual, norm_w, quant, 3, x2=routed) + 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) + + def _swiglu_mlp(self, h: torch.Tensor, prefix: str, push: Optional[tuple] = None): + """A gated MXFP8 MLP (the dense MLP, 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") + 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", 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 fused router with the NVFP4 activation up to 48 rows, else the router logits and + the experts quantizing their bf16 input.""" + core = self.core + fused = _router_fp4(core, i, h) + if fused is not None: + router_logits, x_fp4, x_sf = fused + return _routed_experts(core, i, x_fp4, x_sf, router_logits, do_finalize) + router_logits = _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) + 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", 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") + 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 MSA backend, o_proj. + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s") + out_fp8 = core.main_kv_is_fp8 or core.main_kv_is_nvfp4 + fused = _qk_norm_rope(core, i, qkv, position_ids, out_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", 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") + 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", push) + + 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") + 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", 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 a span of + at most ``FUSED_SCORE_SELECT_MAX_REQUESTS`` requests scores and selects in one launch; a wider span runs the + CuTe DSL scorer (the per-index-head block maxima, reduced to the KV head) and the long-row selector (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 + fused = ( + None if identity else self._fused_score_select(i, idx_q, step, rows, qlen, num_tokens) + ) + if identity: + topk = step["msa_short_table"][: num_kv_heads * num_tokens * MSA_REQUIRED_TOPK].view( + num_kv_heads, num_tokens, MSA_REQUIRED_TOPK + ) + elif fused is not None: + topk = fused.permute(1, 0, 2) + 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 + ) + 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, + ) + 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 + ) + 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=num_tokens <= SELECT_EARLY_MAX_ROWS, + ) + if table is None: + 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) + 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, + ) + + def _fused_score_select(self, i: int, idx_q, step, rows: int, qlen: int, num_tokens: int): + """The top-16 table [total_q, num_kv_heads, 16] of a covered decode span from the fused score + select, or None + (after a miss when the span is covered) for the scorer + selector.""" + core = self.core + path = FUSED_SCORE_SELECT.path + if not FUSED_SCORE_SELECT.covers(rows): + uncovered( + path, f"more than {FUSED_SCORE_SELECT_MAX_REQUESTS} requests: scorer + selector" + ) + return None + config = core.attn[i].m3_config + idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) + idx_k_pages = step["msa_idx_k_cache"](i).squeeze(1) + max_tiles = int(step["msa_worst_case_max_k_tiles"]) + n_valid = step["msa_n_valid_blocks"] + reason = _fused_score_select_kernel.supported( + idx_q_view, idx_k_pages, qlen, config.num_kv_heads, max_tiles + ) + if reason is None and ( + num_tokens != rows * qlen + or n_valid is None + or n_valid.dtype != torch.int32 + or int(n_valid.shape[0]) < num_tokens + ): + reason = "the step's query rows or valid-block counts are not one decode span" + if reason is not None: + missed(path, rows, reason) + return None + bufs = _fused_score_select_kernel.buffers(idx_q.device, max_tiles) + if bufs is None: + capture_missed(path, rows, "the fused score + select scratch") + return None + _record("m3_fused_score_select", rows, qlen) + return _fused_score_select_kernel.fused_score_select( + idx_q_view, + idx_k_pages, + step["msa_block_table"][:rows], + step["msa_seq_lens_cuda"][:rows], + n_valid[:num_tokens].contiguous(), + decode_query_len=qlen, + num_kv_heads=config.num_kv_heads, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + max_tiles=max_tiles, + scratch=bufs[0], + counters=bufs[1], + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# Speculative decoding: one-model Eagle3 at checkpoint precision +# -------------------------------------------------------------------------------------------------------------------- + + +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 never stores the sampler's outputs itself, so + ``begin_step`` and ``sampler_rows_stored`` always report that nothing was staged. + """ + + def __init__(self) -> None: + self.store: Optional[object] = None + self.step_inputs: Optional[tuple] = 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) + + def bind_sampler_store(self, store: object) -> None: + self.store = store + + def begin_step(self, seq_slots: Optional[list], num_rows: int) -> bool: + return False + + 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: + return False + + 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, + ) + + +# (N, K, M, PDL) of every drafter low-M call compiled before the executor's warmup (``_prebuild_draft_lowm``). +_DRAFT_LOWM_BUILT: set = set() + + +def _draft_linear(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """A drafter BF16 linear ``x @ weight^T``: FlashInfer's fixed-tactic low-M kernel on the rows ``DRAFT_LOWM`` + covers (compiled before warmup; a covered call found uncompiled is a miss), else the in-tree BF16 GEMM.""" + n, k = weight.shape + m = int(x.shape[0]) + path = DRAFT_LOWM.path + if x.dim() != 2 or not DRAFT_LOWM.covers(((n, k), m)): + uncovered(path, f"the drafter's {n}x{k} linear past its tactic table's rows") + elif 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(path, "a strided activation (the in-tree GEMM reads it in place)") + elif x.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: + missed(path, m, "the activation or weight is not BF16") + elif (n, k, m, get_env_enable_pdl()) not in _DRAFT_LOWM_BUILT: + missed(path, m, f"the {n}x{k} kernel at {m} rows was not compiled before warmup") + else: + kind, tactic = tactic_for(n, k, m) + out = torch.empty((m, n), dtype=x.dtype, device=x.device) + _record("m3_draft_lowm", (n, k), m) + return m3_draft_lowm(x, weight, out, kind, tactic, get_env_enable_pdl()) + return m3_bf16_linear(x, weight) + + +def _prebuild_draft_lowm(draft: "M3Eagle3Draft") -> int: + """Compile (and run once) every covered row count of the drafter's linears, outside capture, on the current + stream. Returns the number of calls compiled.""" + use_pdl = get_env_enable_pdl() + built = 0 + shapes = {tuple(draft.w[key].shape): draft.w[key] for key in ("qkv", "o", "gu", "dn", "fc")} + for (n, k), weight in sorted(shapes.items()): + for _, m in sorted(key for key in DRAFT_LOWM.keys if key[0] == (n, k)): + if (n, k, m, use_pdl) in _DRAFT_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_BUILT.add((n, k, m, use_pdl)) + built += 1 + return built + + +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: 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) + + @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 _draft_linear(hs, self.w["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 = _draft_linear(x, self.w["qkv"]) + _record("m3_draft_attention", int(md.num_tokens)) + o = m3_draft_attention(self.attn, qkv, md) + out = _draft_linear(o, self.w["o"]) + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def mlp(self, h: torch.Tensor, reduce: bool) -> torch.Tensor: + gu = _draft_linear(h, self.w["gu"]) + act = m3_silu_and_mul(gu, None, None, None) + out = _draft_linear(act, self.w["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. Its BF16 linears run FlashInfer's fixed-tactic low-M + kernels at the rows ``DRAFT_LOWM`` covers (``_draft_linear``). 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 + 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 + 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_hidden_states( + self, spec_metadata, layer_id: int, hidden_states, residual=None + ) -> None: + """The Eagle3 capture with this step's capture norm; 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``).""" + _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: + """The fused draft loop's head copy, the drafter linears' low-M kernels and the drafter attention's KV-split + cap, from the loaded weights, before the executor's warmup, on every rank.""" + 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 = self.prebuild_drafter_linears(draft) + 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 + ) + ) + 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, + ) - def __init__(self, core: MiniMaxM3Tp8ModelingV2) -> None: - self.core = core + 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, + ) - def forward( - self, attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, **kwargs - ): - core = self.core - hidden_states = core.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds - residual = None - for layer in core.layers: - hidden_states, residual = layer( + 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, - residual=residual, spec_metadata=spec_metadata, - **kwargs, + draft_model=draft_model, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, ) - return hidden_states - + else: + # Save the old attn_metadata and spec_metadata + self._prepare_attn_metadata_for_spec_dec(attn_metadata) -class DecodeTarget: - """A pure generation step of at most 48 tokens. - - 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-row engine, pushing its - sum into the boundary all-reduce; the unfinalized MoE with the finalize in - the all-reduce from 9 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. ``dispatch.py`` lists the row - ranges and what a covered step that cannot take its path does. - """ + # 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, + ) - def __init__(self, core: MiniMaxM3Tp8ModelingV2) -> None: - self.core = core - self._moe_push_descs: Dict[tuple, Any] = {} + # 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 - def forward( - self, attn_metadata, input_ids, position_ids, inputs_embeds, spec_metadata, **kwargs - ): - core = self.core - num_tokens = int(input_ids.numel()) - plan = plan_for(num_tokens) - hidden_states, residual = self._embed(input_ids, num_tokens) - for layer in core.layers: - if residual is None or ( - spec_metadata is not None and spec_metadata.is_layer_capture(layer.layer_idx) - ): - hidden_states, residual = layer( - position_ids=position_ids, - hidden_states=hidden_states, - attn_metadata=attn_metadata, - residual=residual, - spec_metadata=spec_metadata, - **kwargs, - ) - else: - hidden_states, residual = self._layer( - layer, plan, position_ids, hidden_states, residual, attn_metadata, kwargs - ) - return hidden_states + # Get the draft KV cache manager if using separate layouts + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - 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 dispatch.EMBED_FIRSTNORM.covers(num_tokens): - dispatch.uncovered( - dispatch.EMBED_FIRSTNORM.path, f"more than {EMBED_FIRSTNORM_MAX_TOKENS} tokens" + 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, ) - return core.embed_tokens(input_ids), None - owner = core._embed_owners.get((num_tokens, input_ids.dtype)) - if owner is None or not input_ids.is_contiguous(): - dispatch.missed( - dispatch.EMBED_FIRSTNORM.path, - num_tokens, - f"no embedding owner for {input_ids.dtype} ids (built for int32 / int64 by build_targets)" - if owner is None - else "the input ids are not contiguous", + # 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, ) - return core.embed_tokens(input_ids), None - return m3_embedding_firstnorm( - input_ids, - owner.table, - owner.norm.weight, - owner.eps, - owner.raw, - owner.output, - owner.use_pdl, - ) - @staticmethod - def _bucket_flags(layer: M3DecoderLayer, num_tokens: int, hidden: int): - """(bucket workspace, its 8-CTA call counters) of the boundary all-reduce for ``num_tokens``, or None.""" - workspace = m3_mc_allreduce.get_workspace(layer.mapping, 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) + attn_metadata.use_spec_decoding = True - @staticmethod - def _no_bucket(path: str, num_tokens: int) -> None: - if torch.cuda.is_current_stream_capturing(): - dispatch.capture_missed(path, num_tokens, "the multicast all-reduce workspace") - else: - dispatch.missed(path, num_tokens, "no multicast all-reduce bucket holds the step") + # 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, + ) - def _attn_push_target(self, layer: M3DecoderLayer, num_tokens: int, hidden: int): - """The o_proj's push tuple into the attention boundary's all-reduce, or None (after a miss). + 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, + } - 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(layer, num_tokens, hidden) - if found is None: - self._no_bucket(dispatch.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 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 _moe_push_targets(self, layer: M3DecoderLayer, num_tokens: int, hidden: int): - """(engine push descriptor, shared down-projection push tuple) for the MoE boundary's exchange sets, or None - (after a miss). + 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 - 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 - (``M3MoE.prebuild_engines``), so it does.""" - from .ops.moe.m3_moe_engine import ar_push_desc + 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. - found = self._bucket_flags(layer, num_tokens, hidden) - if found is None: - self._no_bucket(dispatch.MOE_PUSH.path, num_tokens) + 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 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.""" + 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 - workspace, flags = found - key = (workspace.mc.data_ptr(), flags.data_ptr()) - desc = self._moe_push_descs.get(key) - if desc is 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(): - dispatch.capture_missed( - dispatch.MOE_PUSH.path, num_tokens, "the MoE engine's push descriptor" - ) + capture_missed("spec_tail", None, "the stacked Eagle3 fc_norm weights") 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, + cache = ( + draft_model.w["fc_norm0"], + torch.stack( + [draft_model.w[f"fc_norm{i}"].detach() for i in range(layers)] + ).contiguous(), ) - 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, - ) + 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) + 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_fc_fork( + self, + spec_metadata, + num_tokens: int, + draft_model: Optional[M3Eagle3Draft], + ) -> Optional[torch.cuda.Stream]: + """Launch the first draft pass's fc_norm + fc on a forked stream beside the target head GEMM (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() + 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 + ) + 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 _boundary( - self, layer: M3DecoderLayer, use_mc: bool, x, residual, norm, x2=None, quant=False + 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 + 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, ): - """All-reduce + (second addend) + residual + RMSNorm: the multicast kernel, else the built-in all-reduce.""" - if use_mc: - workspace = m3_mc_allreduce.get_workspace(layer.mapping, x.shape[-1]) - fused = ( - None - if workspace is None - else m3_mc_allreduce.m3_mc_allreduce_norm( - x, - residual.contiguous(), - norm.weight, - norm.variance_epsilon, - workspace, - use_pdl=get_env_enable_pdl(), - x2=x2, - quant=quant, + """Fused verification (see tail_eligible); same outputs as the strict acceptance path.""" + 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, + ) + 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 ) + 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 (none at FC_EARLY_MAX_ROWS = 0), + # else one beside the head. + fork = self._m3_fc_early_stream + self._m3_fc_early_stream = None + if fork is None: + 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 BF16 GEMM, whose all-gather the + # fused kernels replace (they exchange per-row records across ranks, not logits). + logits = lm_head.apply_linear(rows, None) + 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, ) - if fused is not None: - return fused - dispatch.missed( - dispatch.MC_ALLREDUCE.path, - x.shape[0], - m3_mc_allreduce.decline_reason(x, residual, workspace, x2=x2), + 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, ) - if x2 is not None: - x = x.add_(x2) - return layer.native_allreduce_norm(x, residual, norm) - - def _prepushed(self, layer, x, residual, norm, quant, prepushed, x2=None): - workspace = m3_mc_allreduce.get_workspace(layer.mapping, x.shape[-1]) - fused = m3_mc_allreduce.m3_mc_allreduce_norm( - x, - residual.contiguous(), - norm.weight, - norm.variance_epsilon, - workspace, - use_pdl=get_env_enable_pdl(), - x2=x2, - quant=quant, - prepushed=prepushed, - ) - if fused is None: + path = "spec_tail" + 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( - f"layer {layer.layer_idx}: rows were pushed into an all-reduce that did not run" + "M3 fused verify unavailable after tail_eligible() accepted the step" ) - return fused + return result + + def _m3_force_vec( + self, batch_size: int, base_total: int, max_total: int, device + ) -> Optional[torch.Tensor]: + """Per-request forced accepted counts for a fractional forced acceptance, else None. + + Draws exactly as _apply_force_accepted_tokens does for a decode-only batch (same fixed-seed pool, one counter + bump per step, the same (counter, slot) hash), so every TP rank and the generic path agree on the counts. + """ + 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) + 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 + ) + return True + + def prebuild_drafter_linears(self, draft_model: Optional[M3Eagle3Draft]) -> int: + """Compile every drafter linear call ``DRAFT_LOWM`` covers, and cap the drafter attention's trtllm-gen KV + split per request count (``DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH``). Called before the executor's warmup, on every + rank. Returns the number of low-M calls compiled.""" + if draft_model is None: + return 0 + draft_model.attn.dense_decode_max_seq_len_by_batch = DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH + return _prebuild_draft_lowm(draft_model) + + 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_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 _layer( + def _forward_draft_loop( self, - layer: M3DecoderLayer, - plan: DecodePlan, - position_ids, - hidden_states, - residual, + inputs, attn_metadata, - kwargs, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, ): - num_tokens = hidden_states.shape[0] - hidden = hidden_states.shape[-1] - # 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: - dispatch.uncovered( - "decode_graph", - "an eager decode step: the decode-graph GEMV dispatch and the boundary pushes are off", + """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, ) - push = plan.attn_mc_ar and decode_graph and dispatch.ATTN_PUSH.covers(num_tokens) - armed = None - if push: - armed = self._attn_push_target(layer, num_tokens, hidden) - m3_prequant.AR_PUSH.arm(armed) - hidden_states = layer.self_attn( - position_ids=position_ids, - hidden_states=hidden_states, - attn_metadata=attn_metadata, - all_reduce_params=AllReduceParams(enable_allreduce=False), - **kwargs, + 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, ) - if push: - m3_prequant.AR_PUSH.target = None - if push and m3_prequant.AR_PUSH.take_pushed(hidden_states): - hidden_states, residual = self._prepushed( - layer, hidden_states, residual, layer.post_attention_layernorm, False, AR_PUSH_MODE - ) - else: - if armed is not None: - dispatch.missed( - dispatch.ATTN_PUSH.path, - num_tokens, - "the o_proj did not push its rows into the armed boundary all-reduce", - ) - hidden_states, residual = self._boundary( - layer, plan.attn_mc_ar, hidden_states, residual, layer.post_attention_layernorm - ) - moe = layer.block_sparse_moe - if moe is None: - out = layer.mlp( - hidden_states, final_all_reduce_params=AllReduceParams(enable_allreduce=False) - ) - return self._boundary(layer, plan.attn_mc_ar, out, residual, layer.next_layer_layernorm) + 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``): - norm = layer.next_layer_layernorm - quant = plan.qkv_prequant and layer.next_is_decoder_layer - path = dispatch.moe_path(num_tokens) - moe_push = None - if path == dispatch.MOE_ENGINE.path and moe.shared_experts is not None: - # The engine itself runs inside the MoE forward; looked up (or found missing) here first so a push is - # armed only for an engine that will take it. - engine = moe.engine(num_tokens, shared=True) - if engine is not None and plan.moe_mc_ar and decode_graph: - moe_push = self._moe_push_targets(layer, num_tokens, hidden) - if moe_push is not None: - m3_prequant.MOE_PUSH.arm(*moe_push) - if path == dispatch.MC_FINALIZE.path and plan.moe_mc_ar and moe.shared_experts is not None: - if self._bucket_flags(layer, num_tokens, hidden) is None: - self._no_bucket(dispatch.MC_FINALIZE.path, num_tokens) - else: - shared, fc2, scales, perm = moe( - hidden_states, - attn_metadata, - final_all_reduce_params=AllReduceParams(enable_allreduce=False), - do_finalize=False, - ) - 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.weight, - norm.variance_epsilon, - m3_mc_allreduce.get_workspace(layer.mapping, hidden), - use_pdl=get_env_enable_pdl(), - quant=quant, + * 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, ) - if fused is None: + h = draft_model.attention(x, md, reduce=False) + h, residual = draft_model.ar_norm(h, residual, "post_norm") + h = draft_model.mlp(h, reduce=False) + h, to_save = draft_model.ar_norm(h, residual, "final_norm") + 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( - "the MC finalize all-reduce rejected the unfinalized MoE outputs" + "M3 fused draft step unavailable after _m3_fused_loop_ok() accepted the step" ) - return fused - parts = moe( - hidden_states, - attn_metadata, - final_all_reduce_params=AllReduceParams(enable_allreduce=False), - return_parts=True, + 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" ) - shared, routed = parts if isinstance(parts, tuple) else (parts, None) - if moe_push is not None: - m3_prequant.MOE_PUSH.disarm() - if routed is None: - # The engine folded the shared expert in and pushed routed + shared into set A. - folded_pushed = m3_prequant.MOE_PUSH.engine_pushed == shared.data_ptr() - m3_prequant.MOE_PUSH.engine_pushed = m3_prequant.MOE_PUSH.shared_pushed = None - if folded_pushed: - return self._prepushed(layer, shared, residual, norm, quant, 2) - else: - routed_pushed, shared_pushed = m3_prequant.MOE_PUSH.take_pushed(shared, routed) - if routed_pushed and shared_pushed: - return self._prepushed(layer, shared, residual, norm, quant, 3, x2=routed) - if routed_pushed or shared_pushed: - raise RuntimeError("only one MoE part pushed into the boundary all-reduce") - dispatch.missed( - dispatch.MOE_PUSH.path, - num_tokens, - "the MoE did not push its output into the armed boundary all-reduce", - ) - return self._boundary(layer, plan.moe_mc_ar, shared, residual, norm, x2=routed, quant=quant) + 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, measured on GB300 AgentX decode (tp4 and tp8). + + * ``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. + * 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. + """ + return RuntimeHooks( + pack_step_inputs=True, + skip_unchanged_h2d=True, + loop_on_execution_stream=True, + merge_tuned_tactics_across_ranks=True, + skip_greedy_rng_upload=True, + spec_step_fold=spec_fold, + msa=MsaRuntimePolicy( + short_sequence_identity=True, + kv_update=_msa_kv_update_kernel.msa_kv_update, + skip_unchanged_block_offsets=True, + reuse_pool_views=True, + ), + ) -class ModelingV2MiniMaxM3Nvfp4Sm103Tp8(M3SpecDecForwardMixin, MiniMaxM3ForCausalLM): +class ModelingV2MiniMaxM3Nvfp4Sm103Tp8( + DecoderModelForCausalLM[MiniMaxM3Tp8ModelingV2, PretrainedConfig] +): """MiniMax-M3 NVFP4 / sm_103 / tp8, text only, with one-model Eagle3 or no speculation. - A subclass of the built-in text model shell for its loader, its KV-cache and - transceiver preferences and its one-engine speculative plumbing; the text - model it wraps, the forward and the speculative worker are this target's. + 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: @@ -561,64 +4206,132 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: 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 (dispatch.py). - dispatch.set_require(resolved_mode(type(self)) is ModelingV2Mode.REQUIRE) - if is_minimax_m3_vl_config(model_config.pretrained_config): - model_config = get_text_model_config(model_config) + # 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 - # The one-engine shell's constructor, past the built-in one that builds the built-in text model. - SpecDecOneEngineForCausalLM.__init__( - self, MiniMaxM3Tp8ModelingV2(model_config), model_config + super().__init__( + MiniMaxM3Tp8ModelingV2(model_config), + config=model_config, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, ) - self.runtime_hooks, self._spec_fold = make_runtime_hooks(self) - if self.spec_worker is not None: - self.spec_worker = M3Eagle3Worker.adopt(self.spec_worker) - worker = self.spec_worker - for layer in self.model.layers: - layer.capture = worker.capture_hidden_states + 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) + spec_fold = M3SpecStepFold() + self.runtime_hooks = _runtime_hooks(spec_fold) def load_weights(self, weights: Dict, *args, **kwargs) -> None: - _weights.load(self, weights, *args, **kwargs) + _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 setup_aliases(self) -> None: - """Chain each decoder layer's next_layer_layernorm for the fused boundaries.""" - super().setup_aliases() - layers = self.model.layers - for idx, layer in enumerate(layers): - layer.next_is_decoder_layer = idx < len(layers) - 1 + 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 fused draft loop's head copy and the drafter linears' low-M kernels.""" super().cache_derived_state() - self.model.build_targets() + self.model.post_load_weights() if self.spec_worker is not None: - # The fused draft loop's head copy, from the loaded draft weights, before the executor's warmup. - start = time.perf_counter() - if self.spec_worker.prebuild_draft_head(self.draft_model): - torch.cuda.synchronize() - logger.info( - "[M3 prebuild] fused draft head built before warmup: %.1f s" - % (time.perf_counter() - start) - ) - # The drafter linears' low-M kernels and the drafter attention's KV-split cap, before the warmup. - start = time.perf_counter() - built = self.spec_worker.prebuild_drafter_linears(self.draft_model) - torch.cuda.synchronize() - logger.info( - "[M3 prebuild] drafter low-M GEMMs built before warmup: %d calls, %.1f s" - % (built, time.perf_counter() - start) - ) + self.spec_worker.prebuild(self.lm_head, self.draft_model) - def forward(self, attn_metadata: AttentionMetadata, *args, **kwargs): - """The fused-verify one-engine forward (``M3SpecDecForwardMixin``), text only.""" - multimodal_params: List[Any] = kwargs.pop("multimodal_params", None) or [] - if any( - param is not None and getattr(param, "multimodal_data", None) - for param in multimodal_params - ): + 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" ) - return super().forward(attn_metadata, *args, **kwargs) + 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, + **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) + return 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, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/ops/activation/m3_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/activation/m3_silu_and_mul.py new file mode 100644 index 000000000000..7a494d97d788 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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 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_103}) + + 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_103__tp8/ops/attention/m3_draft_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_draft_attention.py new file mode 100644 index 000000000000..6347348d1681 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/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_103}) + + 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_103__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..bcd928fccf5b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp8/ops/attention/m3_msa_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_msa_attention.py new file mode 100644 index 000000000000..60990ce4a415 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/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_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_103}) + + 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_103__tp8/ops/attention/m3_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..681aec29472e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp8/ops/attention/m3_qkv_insert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_qkv_insert.py new file mode 100644 index 000000000000..fc46936cc839 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp8/ops/attention/m3_select_blocks.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_select_blocks.py new file mode 100644 index 000000000000..6bf1764bfc0d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp8/ops/comm/m3_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_allreduce.py new file mode 100644 index 000000000000..7e03223a5dba --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + # 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_103__tp8/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/ops/embedding/m3_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding/m3_embedding.py new file mode 100644 index 000000000000..f979fdecf376 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp8/ops/gemm/m3_bf16_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_bf16_linear.py new file mode 100644 index 000000000000..5007e1da7386 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp8/ops/gemm/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_draft_lowm.py index 4c3eca583b61..1caeda1aa98d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_draft_lowm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_draft_lowm.py @@ -54,22 +54,14 @@ def tactic_for(n: int, k: int, m: int) -> Optional[Tuple[str, tuple]]: return None -def _cells() -> tuple[Cell, ...]: - names = { - (1280, 12288): "qkv", - (6144, 1024): "o_proj", - (4608, 6144): "gate_up", - (6144, 2304): "down", - (6144, 18432): "fc", - } - return tuple( - Cell( - why=f"drafter {names[nk]}, {m} rows: {tactic_for(nk[0], nk[1], m)[0]} tactic", - spec=dict(nk=nk, m=m), - ) - for nk, table in TACTICS.items() - for m in range(1, table[-1][0] + 1) - ) +#: The drafter linear each (N, K) of the table is. +_NAMES: Mapping[Tuple[int, int], str] = { + (1280, 12288): "qkv", + (6144, 1024): "o_proj", + (4608, 6144): "gate_up", + (6144, 2304): "down", + (6144, 18432): "fc", +} class M3DraftLowM(OpWrapper): @@ -78,11 +70,19 @@ class M3DraftLowM(OpWrapper): ``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 (kernels/m3_draft_lowm.prebuild). + covered call before warmup (modeling's ``_prebuild_draft_lowm``). """ ARCHS = frozenset({Arch.SM_103}) - CELLS: tuple[Cell, ...] = _cells() + # Every row count of every table entry. + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"drafter {_NAMES[nk]}, {m} rows: {tactic_for(nk[0], nk[1], m)[0]} tactic", + spec=dict(nk=nk, m=m), + ) + for nk, table in TACTICS.items() + for m in range(1, table[-1][0] + 1) + ) note = """ The table is the TP8 drafter's: its shapes are the per-rank shards of the Eagle3 layer at 8-way tensor diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_fp32_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_fp32_linear.py new file mode 100644 index 000000000000..8d7374810675 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp8/ops/gemm/m3_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mxfp8_gemm.py new file mode 100644 index 000000000000..8fc131a7bb61 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/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_103}) + + # 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)), + ) + + 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_103__tp8/ops/gemm/m3_mxfp8_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mxfp8_linear.py new file mode 100644 index 000000000000..1bdc1829508d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103}) + + 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_103__tp8/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/m3_embedding_firstnorm.py index ff7fb16e2266..535f928018f5 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/m3_embedding_firstnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/m3_embedding_firstnorm.py @@ -9,10 +9,8 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.embedding_firstnorm.op import ( - DecodeEmbedding, - m3_embed_norm, -) + +from ...kernels.embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm __all__ = ["DecodeEmbedding", "m3_embedding_firstnorm"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_force_vec.py new file mode 100644 index 000000000000..1905516a8403 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/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 ...kernels.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_103}) + + 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_103__tp8/runtime.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/runtime.py deleted file mode 100644 index cfcff73e3cea..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/runtime.py +++ /dev/null @@ -1,50 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The step-loop hooks this target opts into (``RuntimeHooks``), measured on GB300 AgentX decode (tp4 and tp8). - -* ``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. -* 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. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Optional, Tuple - -from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy, RuntimeHooks - -from .kernels import msa_kv_update -from .spec_fold import M3SpecStepFold - -if TYPE_CHECKING: - from .modeling import ModelingV2MiniMaxM3Nvfp4Sm103Tp8 - - -def make_runtime_hooks( - model: "ModelingV2MiniMaxM3Nvfp4Sm103Tp8", -) -> Tuple[RuntimeHooks, Optional[M3SpecStepFold]]: - """This target's hooks, and the speculative step fold they carry (None without speculation).""" - spec_fold = M3SpecStepFold() if model.spec_config is not None else None - hooks = RuntimeHooks( - pack_step_inputs=True, - skip_unchanged_h2d=True, - loop_on_execution_stream=True, - merge_tuned_tactics_across_ranks=True, - skip_greedy_rng_upload=True, - spec_step_fold=spec_fold, - msa=MsaRuntimePolicy( - short_sequence_identity=True, - kv_update=msa_kv_update.msa_kv_update, - skip_unchanged_block_offsets=True, - reuse_pool_views=True, - ), - ) - return hooks, spec_fold diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec_fold.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec_fold.py deleted file mode 100644 index dbc210c0ab46..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec_fold.py +++ /dev/null @@ -1,112 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""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 never stores the sampler's outputs itself, so ``begin_step`` and -``sampler_rows_stored`` always report that nothing was staged. - -* ``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``). -""" - -from __future__ import annotations - -from typing import Optional - -import torch - -from .kernels import m3_spec_tail - - -class M3SpecStepFold: - """``SpecStepFold`` implementation of the MiniMax-M3 tp8 target (see the module docstring). - - One instance per target model; the target passes it to ``RuntimeHooks.spec_step_fold``. - """ - - def __init__(self) -> None: - self.store: Optional[object] = None - self.step_inputs: Optional[ - tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] - ] = 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) - - def bind_sampler_store(self, store: object) -> None: - self.store = store - - def begin_step(self, seq_slots: Optional[list[int]], num_rows: int) -> bool: - return False - - 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: - return False - - 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, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec_worker.py deleted file mode 100644 index ce1f5bf51968..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec_worker.py +++ /dev/null @@ -1,1119 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Eagle3 one-model speculative decoding of the MiniMax-M3 tp8 target, at stock numerics. - -``M3Eagle3Worker`` is the in-tree ``Eagle3OneModelWorker`` with a fused verify and a fused draft loop: - -* Verify: on decode-only batches whose verify rows the fused kernels cover, the target's BF16 lm_head computes this - rank's vocabulary shard of the logits (the in-tree LM 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 in-tree Eagle3 drafter's modules and weights at their precision (BF16 weights and activations, the - full-vocabulary BF16 head, greedy drafts). The first pass may take the per-layer fc_norm RMSNorms in one launch (or - folded into the hidden-state capture) and run its fc beside the target head GEMM. Its BF16 linears run - FlashInfer's fixed-tactic low-M kernels at the rows ``dispatch.DRAFT_LOWM`` covers (``prebuild_drafter_linears``). - On a decode-only step of at most 8 requests (``dispatch.DRAFT_STEP``) the draft loop is fused - (``_m3_fused_draft_loop``): each pass's draft input in one launch, the drafter layer's in-tree 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. - -The target's outer CausalLM drives the worker through ``M3SpecDecForwardMixin.forward`` (or the same calls made from -its own forward), and its layer loop captures the Eagle3 hidden states through ``capture_hidden_states``. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any, Optional - -import torch -import triton -import triton.language as tl -from torch import nn - -from tensorrt_llm._torch.distributed import AllReduce, AllReduceParams -from tensorrt_llm._torch.models.modeling_speculative import _slice_spec_position_ids -from tensorrt_llm._torch.modules.gated_mlp import GatedMLP -from tensorrt_llm._torch.modules.linear import TensorParallelMode -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager -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, -) -from tensorrt_llm.mapping import Mapping - -from . import dispatch -from .kernels import m3_draft_lowm, m3_draft_step, m3_sampled_accept, m3_shard_accept, m3_spec_tail - -if TYPE_CHECKING: - from tensorrt_llm._torch.model_config import ModelConfig - from tensorrt_llm.llmapi.llm_args import EagleDecodingConfig - -# Id of the target's final hidden state in ``layers_to_capture`` (SpecDecOneEngineForCausalLM.layer_idx). -FINAL_HIDDEN_LAYER_ID = -1 - - -@triton.jit -def _m3_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) - - -def draft_ar_norm( - all_reduce: AllReduce, x: torch.Tensor, residual: torch.Tensor, norm: nn.Module -) -> tuple[torch.Tensor, torch.Tensor]: - """One draft-layer boundary in stock Eagle3's structure: the projection's plain all-reduce (BF16 out), then the - norm module with the residual (FlashInfer's fused_add_rmsnorm: FP32 add, norm of the FP32 sum). Returns the - normalized rows and the residual sum.""" - return norm(all_reduce(x), residual) - - -def capture_hidden_states( - spec_metadata: Any, - layer_id: int, - hidden_states: torch.Tensor, - residual: Optional[torch.Tensor], - capture_norm: Optional[tuple[torch.Tensor, float]], -) -> 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) (``M3Eagle3Worker.capture_norm``), 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 " - f"buffer capacity: num_tokens={num_tokens}, " - f"capacity={spec_metadata.hidden_states.shape[0]}" - ) - if num_tokens > hidden_states.shape[0]: - raise RuntimeError( - "EAGLE3 hidden-state capture token count exceeds " - f"available hidden states: num_tokens={num_tokens}, " - f"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 - - -class M3Eagle3Worker(Eagle3OneModelWorker): - """``Eagle3OneModelWorker`` with the MiniMax-M3 fused verify (see the module docstring). - - Methods whose docstring says "Mirrors" are copies of the named in-tree method with the fused-verify branches - inserted; compare them with that method when it changes. - """ - - def __init__( - self, - spec_config: "EagleDecodingConfig", - mapping: Optional[Mapping] = None, - model_config: Optional["ModelConfig"] = None, - use_separate_draft_kv_cache: bool = False, - ): - super().__init__( - spec_config, - mapping=mapping, - model_config=model_config, - use_separate_draft_kv_cache=use_separate_draft_kv_cache, - ) - self._init_fused_verify() - - @classmethod - def adopt(cls, worker: Eagle3OneModelWorker) -> "M3Eagle3Worker": - """Turn the in-tree worker the one-engine shell built (``get_spec_worker``) into this worker, in place.""" - if type(worker) is not Eagle3OneModelWorker: - raise TypeError(f"expected an Eagle3OneModelWorker, got {type(worker).__name__}") - worker.__class__ = cls - worker._init_fused_verify() - return worker - - def _init_fused_verify(self) -> None: - # 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_norm[0].weight, stacked fc_norm weights [layers, H]), built eagerly on the first fused first pass. - self._m3_fc_norm_w: Optional[tuple[torch.Tensor, torch.Tensor]] = None - # Armed per step by arm_capture_norm: (stacked fc_norm weights, eps) or None. - self._m3_capture_norm: Optional[tuple[torch.Tensor, float]] = 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 - # (capture_hidden_states), and the fork that did it, joined by _m3_target_accept. - self._m3_fc_early_drafter: Optional[nn.Module] = None - self._m3_fc_early_stream: Optional[torch.cuda.Stream] = None - self._m3_force_vec_warm = False - # The fragment-ordered copy of the drafter's BF16 head shard the fused draft loop streams, built before the - # executor's warmup (prebuild_draft_head). - self._m3_draft_head: Optional[m3_draft_step.DraftHead] = 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 - - @property - def capture_norm(self) -> Optional[tuple[torch.Tensor, float]]: - """What ``capture_hidden_states`` needs for this step's captures (set by ``arm_capture_norm``).""" - return self._m3_capture_norm - - def capture_hidden_states( - self, - spec_metadata: Any, - layer_id: int, - hidden_states: torch.Tensor, - residual: Optional[torch.Tensor] = None, - ) -> None: - """``capture_hidden_states`` with this step's capture norm; 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``).""" - capture_hidden_states( - 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 _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 - ) - ) - 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 because - # the accepted token count affects which Mamba states are valid. The - # isinstance check below naturally no-ops on non-Mamba kv_cache_managers, - # so this is safe to run unconditionally regardless of spec mode (Eagle3 - # over a Mamba-style draft is plausible, even if no such draft exists today). - 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: torch.LongTensor, - position_ids: torch.LongTensor, - hidden_states: torch.Tensor, - accepted_tokens: torch.Tensor, - attn_metadata: Any, - spec_metadata: Any, - draft_model: nn.Module, - ): - """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: Any, spec_metadata: Any) -> 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) > dispatch.DECODE_MAX_TOKENS: - return False - return not spec_metadata.wants_advanced_draft_sampling - - def tail_eligible(self, attn_metadata: Any, spec_metadata: Any, 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 ``dispatch.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 self._m3_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 dispatch.SHARD_ACCEPT.covers(key): - dispatch.uncovered( - dispatch.SHARD_ACCEPT.path, - f"draft length {draft_len} 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: Any, spec_metadata: Any, draft_model: Optional[nn.Module] - ) -> 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.""" - return self._m3_structural_ok(attn_metadata, spec_metadata) and self._m3_fc_norm_ok( - draft_model - ) - - def _m3_seed_first_pass( - self, num_accepted_tokens: torch.Tensor, spec_metadata: Any, 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: torch.Tensor, - accepted_tokens: torch.Tensor, - num_tokens: int, - num_contexts: int, - spec_metadata: Any, - draft_model: nn.Module, - ) -> Optional[dict[str, Any]]: - """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 in-tree 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 not self._m3_fc_norm_ok(draft_model) - ): - return None - dm = draft_model.model - layers = len(spec_metadata.layers_to_capture) - if dm.hidden_size_in != spec_metadata.hidden_size or layers != len(dm.fc_norm): - return None - cache = self._m3_fc_norm_w - if cache is None or cache[0] is not dm.fc_norm[0].weight: - if torch.cuda.is_current_stream_capturing(): - dispatch.capture_missed("spec_tail", None, "the stacked Eagle3 fc_norm weights") - return None - cache = ( - dm.fc_norm[0].weight, - torch.stack([n.weight.detach() for n in dm.fc_norm]).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], dm.fc_norm[0].variance_epsilon) - hidden = dm.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} - - @staticmethod - def _m3_fc_norm_ok(draft_model: Optional[nn.Module]) -> bool: - """An Eagle3 drafter with an fc over several captured layers, each normalized by its own fc_norm (one - shared epsilon): what the fused first drafter inputs compute.""" - dm = getattr(draft_model, "model", None) - fc_norm = getattr(dm, "fc_norm", None) - return ( - fc_norm is not None - and getattr(dm, "fc", None) is not None - and len({n.variance_epsilon for n in fc_norm}) == 1 - ) - - def _m3_fc_fork( - self, - spec_metadata: Any, - num_tokens: int, - draft_model: Optional[nn.Module], - ) -> Optional[torch.cuda.Stream]: - """Launch the first draft pass's fc_norm + fc on a forked stream beside the target head GEMM (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 or not self._m3_fc_norm_ok(draft_model): - return None - dm = draft_model.model - layers = len(spec_metadata.layers_to_capture) - if ( - dm.hidden_size_in != spec_metadata.hidden_size - or layers != len(dm.fc_norm) - 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() - 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], dm.fc_norm[0].variance_epsilon, use_pdl=False - ) - self._m3_fc_hidden = dm.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 or - torch.compile, 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: Any, - spec_metadata: Any, - lm_head: nn.Module, - draft_model: Optional[nn.Module], - ) -> 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 - ``dispatch.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 - # 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) - <= dispatch.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 self._m3_fc_norm_ok(draft_model) - and fc_norm_w[0] is draft_model.model.fc_norm[0].weight - and len(spec_metadata.layers_to_capture) == fc_norm_w[1].shape[0] - ): - self._m3_capture_norm = (fc_norm_w[1], draft_model.model.fc_norm[0].variance_epsilon) - - @staticmethod - def _m3_sampled_verify_ok(spec_metadata: Any) -> 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: torch.Tensor, - spec_metadata: Any, - batch_size: int, - draft_len: int, - draft_model: Optional[nn.Module] = None, - ): - """Fused verification (see tail_eligible); same outputs as the strict acceptance path.""" - 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_kernel[(1,)]( - 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=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 - ) - 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 (none at FC_EARLY_MAX_ROWS = 0), - # else one beside the head. - fork = self._m3_fc_early_stream - self._m3_fc_early_stream = None - if fork is None: - fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model) - # This rank's vocabulary shard of the target logits: the in-tree LM head's own local BF16 GEMM, whose - # all-gather the fused kernels replace (they exchange per-row records across ranks, not logits). - logits = lm_head.apply_linear(rows, None) - vocab_offset = self.mapping.tp_rank * lm_head.weight.shape[0] - if spec_metadata.is_all_greedy_sample: - result = m3_shard_accept.shard_accept( - logits, - vocab_offset, - self.mapping, - draft_tokens, - batch_idx, - force_total, - self._m3_next_new, - force_vec=force_vec, - ) - else: - n_rows = rows.shape[0] - seeds, offsets = self._rng_state_per_token(spec_metadata, n_rows) - result = m3_sampled_accept.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, - ) - if fork is not None: - self._m3_fc_join(fork) - if result is None: - 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: torch.device - ) -> Optional[torch.Tensor]: - """Per-request forced accepted counts for a fractional forced acceptance, else None. - - Draws exactly as _apply_force_accepted_tokens does for a decode-only batch (same fixed-seed pool, one counter - bump per step, the same (counter, slot) hash), so every TP rank and the generic path agree on the counts. - """ - 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) - # One launch: bump the shared counter and draw, with _apply_force_accepted_tokens' exact arithmetic. - _m3_force_vec_kernel[(1,)]( - 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=max(16, 1 << (batch_size - 1).bit_length()), - ) - return out - - @staticmethod - def _m3_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 - ) - - # ---------------------------------------------------------------------------------------------- fused draft loop - - def prebuild_draft_head(self, draft_model: Optional[nn.Module]) -> bool: - """Build the fragment-ordered copy of the drafter's BF16 head shard that the fused draft loop streams (rebuilt - from the current weight values on every call). Called by the target's ``cache_derived_state`` once the weights - are real, before the executor's warmup, 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.prepare_head(draft_model.lm_head, self.mapping.tp_rank) - return True - - def prebuild_drafter_linears(self, draft_model: Optional[nn.Module]) -> int: - """Put the drafter's BF16 linears on the fixed-tactic low-M kernels (``dispatch.DRAFT_LOWM``) and compile every - covered call, and cap the drafter attention's trtllm-gen KV split per request count - (``dispatch.DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH``). Called by the target's ``cache_derived_state`` once the - weights are real, before the executor's warmup, on every rank. Returns the number of low-M calls compiled.""" - if draft_model is None: - return 0 - dm = getattr(draft_model, "model", None) - attn = getattr(getattr(getattr(dm, "midlayer", None), "self_attn", None), "attn", None) - if attn is not None: - attn.dense_decode_max_seq_len_by_batch = dispatch.DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH - m3_draft_lowm.install(draft_model) - return m3_draft_lowm.prebuild(draft_model) - - def _m3_drafter_supported(self, draft_model: nn.Module) -> bool: - """The drafter the fused draft loop runs: one plain (non-MLA) Eagle3 layer with a gated MLP over the - target-vocabulary embedding (whole table on every rank), a draft vocabulary equal to the target's (no d2t), - all-reduced TP projections, and a bias-free vocab-sharded BF16 head without padding.""" - dm = getattr(draft_model, "model", None) - if ( - dm is None - or getattr(dm, "num_layers", 0) != 1 - or dm._use_mla - or dm._eh_proj_before_attn - ): - return False - layer = dm.midlayer - if layer._next_layer_regular or not isinstance(layer.mlp, GatedMLP): - return False - embed = dm.embed_tokens - if embed is None or getattr(embed, "tp_size", 1) != 1: - return False - if self.mapping is None or self.mapping.tp_size <= 1 or self.mapping.enable_attention_dp: - return False - if getattr(layer.self_attn.o_proj, "all_reduce", None) is None: - return False - if getattr(layer.mlp.down_proj, "all_reduce", None) is None: - return False - if self._d2t is not None or getattr(dm, "d2t", None) is not None: - return False - lm_head = draft_model.lm_head - return self._m3_sharded_head(lm_head) and m3_draft_step.head_supported(lm_head) - - def _m3_fused_loop_ok( - self, - attn_metadata: Any, - spec_metadata: Any, - batch_size: int, - draft_model: Optional[nn.Module], - ) -> 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) - ``dispatch.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 = dispatch.DRAFT_STEP.path - if not self._m3_drafter_supported(draft_model) or not hasattr( - attn_metadata, "kv_lens_cuda" - ): - dispatch.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 dispatch.DRAFT_STEP.covers((draft_len, batch_size)): - dispatch.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: - dispatch.missed( - path, - batch_size, - "the drafter head's fragment-ordered copy is not built for its current weight (prebuild_draft_head)", - ) - return False - return m3_draft_step.workspace(self.mapping, head.wf.device) is not None - - 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: dict[str, Any], - attn_metadata: Any, - spec_metadata: Any, - draft_model: nn.Module, - draft_kv_cache_manager: Any, - batch_size: int, - num_accepted_tokens: torch.Tensor, - ) -> torch.Tensor: - """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_layernorm 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 in-tree attention and MLP, each followed by the in-tree all-reduce with the residual add - and the next RMSNorm fused into it (post_attention_layernorm, 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 - dm = draft_model.model - layer = dm.midlayer - attn = layer.self_attn - mlp = layer.mlp - attn_all_reduce = attn.o_proj.all_reduce - mlp_all_reduce = mlp.down_proj.all_reduce - table = dm.embed_tokens.weight - 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:] - - no_all_reduce = AllReduceParams(enable_allreduce=False) - 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, - layer.input_layernorm.weight, - layer.input_layernorm.variance_epsilon, - layer.hidden_norm.weight, - layer.hidden_norm.variance_epsilon, - ) - h = attn( - position_ids=pos, - hidden_states=x, - attn_metadata=md, - all_reduce_params=no_all_reduce, - ) - h, residual = draft_ar_norm( - attn_all_reduce, h, residual, layer.post_attention_layernorm - ) - h = mlp(h, final_all_reduce_params=no_all_reduce) - h, to_save = draft_ar_norm(mlp_all_reduce, h, residual, dm.norm) - if dm._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.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 - - -class M3SpecDecForwardMixin: - """The target CausalLM's step forward with the fused-verify hooks of the speculative wrapper. - - Mirrors ``SpecDecOneEngineForCausalLM.forward`` (tensorrt_llm/_torch/models/modeling_speculative.py) with the - fused-verify branches: before the target forward the worker arms the capture norm, and a step the worker verifies - from the logits shard (``M3Eagle3Worker.tail_eligible``) skips the gathered target logits. - - The class it is mixed into (listed before ``DecoderModelForCausalLM``) provides ``model`` (called like the in-tree - decoder model, capturing through ``M3Eagle3Worker.capture_hidden_states``), ``lm_head``, ``logits_processor``, - ``model_config``, ``draft_model`` and ``spec_worker`` (an ``M3Eagle3Worker``, or None without speculation). - """ - - def forward( - self, - attn_metadata: Any, - 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[Any] = None, - resource_manager: Optional[Any] = None, - **kwargs, - ) -> Any: - spec_worker: Optional[M3Eagle3Worker] = self.spec_worker - if spec_worker is not None and spec_metadata is not None: - spec_worker.arm_capture_norm( - attn_metadata, spec_metadata, self.lm_head, self.draft_model - ) - hidden_states = self.model( - input_ids=input_ids, - attn_metadata=attn_metadata, - position_ids=position_ids, - inputs_embeds=inputs_embeds, - spec_metadata=spec_metadata, - **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): - if spec_worker is not None: - spec_worker.capture_hidden_states( - spec_metadata, FINAL_HIDDEN_LAYER_ID, hidden_states, None - ) - else: - capture_hidden_states( - spec_metadata, FINAL_HIDDEN_LAYER_ID, hidden_states, None, None - ) - if attn_metadata.padded_num_tokens is not None: - hidden_states = hidden_states[: attn_metadata.num_tokens] - - if spec_worker is not None: - # The fused verify computes this rank's logits shard itself and needs no gathered logits. - if spec_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 - ) - return spec_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, - ) - - return self.logits_processor.forward( - hidden_states, - self.lm_head, - attn_metadata, - return_context_logits, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/weights.py index c69fa3be8ef1..011445939d14 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/weights.py @@ -1,62 +1,883 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Weight loading: MiniMax-M3 NVFP4 / sm_103 / tp8. - -The checkpoint is the vision-language wrapper's. ``language_model.*`` holds the -language model; ``vision_tower.*``, ``multi_modal_projector.*`` and -``patch_merge_mlp.*`` hold the vision tower and its projector. This target is -text only, so the keys split three ways: - -* ``language_model.*`` goes to the built-in MiniMax-M3 text loader with the - prefix stripped (its multi-token-prediction head, ``model.mtp.*``, is not - loaded, as in the built-in model). That loader packs the five-way - Q/K/V/index-Q/index-K projection, folds the Gemma boundary norms and shards - every tensor for tp8; the modules it fills are this target's subclasses of - the built-in ones, under the built-in names. -* The vision tower and the projector are a predicted non-load: listed here, - never read. -* Any other key fails the load, naming it, rather than being dropped. +"""Weight tables and loaders: MiniMax-M3 NVFP4 / sm_103 / 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 typing import Any, Dict, Optional +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any, Callable -from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper -from tensorrt_llm._torch.models.modeling_minimaxm3 import ( - MiniMaxM3ForCausalLM, - _strip_language_model_prefix, +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, ) -_LANG_PREFIX = "language_model." +#: ``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 -# The vision tower's and the projector's key families: in the checkpoint, not loaded by this text-only target. -PREDICTED_NON_LOAD = ("vision_tower.", "multi_modal_projector.", "patch_merge_mlp.") +def _f32(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return t.to(torch.float32) -def load( - model, - weights: Dict[str, Any], - weight_mapper: Optional[BaseWeightMapper] = None, - params_map: Optional[Dict[str, str]] = None, - allow_partial_loading: bool = False, -) -> None: - """Load the language model's weights into ``model``, the target shell, and check every other key is predicted.""" - unknown = sorted( - k - for k in weights.keys() - if not k.startswith(_LANG_PREFIX) and not k.startswith(PREDICTED_NON_LOAD) + +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 head (rank r holds head r // 2) 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 ) - assert not unknown, ( - f"{len(unknown)} checkpoint key(s) are neither language-model weights nor a predicted non-load, " - f"first {unknown[:5]}" + 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 ) - text_weights, _ = _strip_language_model_prefix( - {k: v for k, v in weights.items() if k.startswith(_LANG_PREFIX)} + 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 ) - assert text_weights, f"the checkpoint has no {_LANG_PREFIX}* keys" - MiniMaxM3ForCausalLM.load_weights( - model, - text_weights, - weight_mapper=weight_mapper, - params_map=params_map, - allow_partial_loading=allow_partial_loading, + 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_103 / 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, each KV head on two + 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 c6b166dc1300..b46be096ad2f 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 @@ -7,6 +7,10 @@ fits it raises ``NoTarget`` stating those facts. In ``auto`` the engine then uses the built-in MiniMax-M3 implementation; in ``require`` it raises that error. +The sparse-attention runtime is part of the decision too (``ROUTES_ON_RUNTIME``): +the targets are written for the MSA implementation with an FP8 indexer and an +FP8 or NVFP4 main KV cache, and any other runtime is the built-in model's. + The checkpoint declares the vision-language wrapper, and the language model's shape lives in its ``text_config``, so that is what this tree reads. The targets are text-only: they load no vision tower and refuse image input. @@ -14,12 +18,17 @@ from __future__ import annotations +import json +from pathlib import Path from typing import Any, Optional, Tuple from tensorrt_llm.llmapi.llm_args import EagleDecodingConfig from tensorrt_llm.quantization.mode import QuantAlgo -from ..._router_index import NoTarget +from ..._router_index import NoTarget, RuntimeFacts + +# ``route`` also reads the engine's sparse-attention runtime and KV cache dtype (``RuntimeFacts``). +ROUTES_ON_RUNTIME = True # The one GPU architecture these targets are written for: the decode kernels # (the tcgen05 GEMVs, the fused MoE engines, the attention megakernels and the @@ -34,6 +43,17 @@ # counts fix the per-rank operand shapes the decode kernels are compiled for. _SHAPE = (60, 6144, 128, 3072, 64, 4) +# The Eagle3 drafter geometry the targets' drafter kernels and weight tables are built for (the GQA drafter: the +# target's hidden size, heads, KV heads and vocabulary, a 18432-wide MLP). +_DRAFTER_GEOMETRY = dict( + hidden_size=6144, + num_attention_heads=64, + num_key_value_heads=4, + head_dim=128, + intermediate_size=18432, + vocab_size=200064, +) + def _field(config: Any, name: str) -> Any: """Read ``name`` off a sub-config that may be an object or a plain dict.""" @@ -94,11 +114,101 @@ def _eagle3_one_model_or_none(spec_config: Optional[Any]) -> bool: ) +def _drafter_mismatch(spec_config: Optional[Any]) -> Optional[str]: + """Why the speculative drafter is not the one the targets' flat drafter builds, or None when it is (or there is + no speculation). + + The targets draft with their own one-layer Eagle3 drafter (``M3Eagle3Draft``): one Llama-style layer, per-layer + fc_norms over several captured target layers, no eh_proj before attention or norm before the fc, no MLA, the + attention and MLP geometry of ``_DRAFTER_GEOMETRY``, the target's hidden size and vocabulary, no biases and an + untied head. Any other Eagle3 drafter is the built-in + implementation's. Read from the drafter checkpoint's ``config.json`` (the engine has resolved + ``speculative_model`` to a local directory before it builds the model). + """ + if spec_config is None: + return None + path = getattr(spec_config, "speculative_model", None) + if path is None: + return "no Eagle3 drafter checkpoint (speculative_model)" + try: + with open(Path(path) / "config.json") as f: + cfg = json.load(f) + except (OSError, ValueError) as err: + return f"the Eagle3 drafter's config.json is not readable ({err})" + eagle = cfg.get("eagle_config") or {} + hidden = cfg.get("hidden_size") + vocab = cfg.get("vocab_size") + captured = getattr(spec_config, "eagle3_layers_to_capture", None) + checks = ( + ( + cfg.get("num_hidden_layers") == 1, + f"{cfg.get('num_hidden_layers')} drafter layers, not 1", + ), + (bool(cfg.get("fc_norm", False)), "no per-layer fc_norm"), + (not eagle.get("eh_proj_before_attn", False), "eh_proj before attention"), + (not eagle.get("norm_before_fc", False), "norm before the fc"), + (not cfg.get("kv_lora_rank"), "an MLA drafter layer"), + ( + cfg.get("target_hidden_size", hidden) == hidden, + "a target hidden size other than the drafter's", + ), + ( + cfg.get("draft_vocab_size", vocab) in (None, vocab), + "a draft vocabulary other than the target's", + ), + ( + not (cfg.get("attention_bias") or cfg.get("mlp_bias") or cfg.get("bias")), + "biased projections", + ), + (not cfg.get("tie_word_embeddings", False), "a tied head"), + (captured is None or len(captured) > 1, "a single captured target layer"), + ) + heads = cfg.get("num_attention_heads") + defaults = dict( + num_key_value_heads=heads, + head_dim=hidden // heads if hidden and heads else None, + ) + geometry = [] + for name, want in _DRAFTER_GEOMETRY.items(): + value = cfg.get(name) + if value is None: + value = defaults.get(name) + geometry.append((value == want, f"{name}={value}, not {want}")) + reasons = [why for ok, why in (*checks, *geometry) if not ok] + return "; ".join(reasons) if reasons else None + + +def _runtime_mismatch(runtime: RuntimeFacts) -> Optional[str]: + """Why the engine's runtime is not the one the targets are written for, or None when it is. + + The MiniMax-M3 MSA backend (it owns the step metadata and the index-K cache), an FP8 indexer and an FP8 or + NVFP4 main KV cache, with the fused five-way QKV/index projection or the separate projections. Every other + runtime the engine accepts (the Triton implementation, the BF16 indexer, the BF16 main KV cache) is the built-in + implementation's. The targets' ``_validate_runtime`` checks the same at construction. + """ + sparse = runtime.sparse_attention_config + if sparse is None or _field(sparse, "algorithm") != "minimax_m3": + return "no MiniMax-M3 sparse_attention_config" + reasons = [] + implementation = _field(sparse, "implementation") + if implementation != "msa": + reasons.append(f"sparse attention implementation {implementation!r}, not 'msa'") + indexer = _field(sparse, "indexer_kv_dtype") + if indexer != "fp8": + reasons.append(f"indexer_kv_dtype {indexer!r}, not 'fp8'") + kv = runtime.kv_cache_quant_algo + if kv not in (QuantAlgo.FP8, QuantAlgo.NVFP4): + reasons.append(f"a {kv if kv is not None else 'BF16'} main KV cache, not FP8 or NVFP4") + return "; ".join(reasons) if reasons else None + + def route( pretrained_config: Any, mapping: Any, compute_capability: Tuple[int, int], spec_config: Optional[Any], + *, + runtime: RuntimeFacts, ) -> type: """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. @@ -124,6 +234,8 @@ def route( quant_algo = _quant_algo(pretrained_config) parallel = _parallel(mapping) speculation = _eagle3_one_model_or_none(spec_config) + drafter = _drafter_mismatch(spec_config) if speculation else None + runtime_miss = _runtime_mismatch(runtime) # MiniMax-M3 NVFP4 on 4 GB300s (sm103), tensor-parallel attention and # experts, one-model Eagle3 or no speculative decoding. @@ -133,6 +245,8 @@ def route( and quant_algo == QuantAlgo.MIXED_PRECISION and parallel == "tp4" and speculation + and drafter is None + and runtime_miss is None ): from .minimax_m3_nvfp4__sm_103__tp4.modeling import ModelingV2MiniMaxM3Nvfp4Sm103Tp4 @@ -146,6 +260,8 @@ def route( and quant_algo == QuantAlgo.MIXED_PRECISION and parallel == "tp8" and speculation + and drafter is None + and runtime_miss is None ): from .minimax_m3_nvfp4__sm_103__tp8.modeling import ModelingV2MiniMaxM3Nvfp4Sm103Tp8 @@ -157,7 +273,11 @@ def route( f"ws={mapping.world_size} tp={mapping.tp_size} pp={mapping.pp_size} cp={mapping.cp_size} " f"moe_tp={mapping.moe_tp_size} moe_ep={mapping.moe_ep_size} " f"attention_dp={mapping.enable_attention_dp}, speculative decoding " - f"{spec_config.spec_dec_mode.name if spec_config is not None else 'off'}; the " + f"{spec_config.spec_dec_mode.name if spec_config is not None else 'off'}" + f"{f' (Eagle3 drafter: {drafter})' if drafter else ''}" + f"{f', runtime: {runtime_miss}' if runtime_miss else ''}; 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" + "tensor-parallel attention and experts, with one-model Eagle3 (a one-layer drafter with per-layer fc_norms) " + "or no speculative decoding, on the MSA sparse attention runtime with an FP8 indexer and an FP8 or NVFP4 KV " + "cache" ) diff --git a/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py b/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py new file mode 100644 index 000000000000..a5c777b5cf7c --- /dev/null +++ b/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.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. +"""Whole-model gate for the MiniMax-M3 modeling_v2 targets. + +Every test here passes ``modeling_v2="require"`` to ``LLM``: under ``"auto"`` a configuration that missed a target's +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). +""" + +import json + +import pytest + +from tensorrt_llm import LLM +from tensorrt_llm._torch._experimental.modeling_v2 import MODELING_V2_VALIDATE_ENV +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.llmapi import ( + CudaGraphConfig, + Eagle3DecodingConfig, + KvCacheConfig, + MiniMaxM3SparseAttentionConfig, + MoeConfig, + SamplingParams, +) +from tensorrt_llm.quantization import QuantAlgo + +from ...conftest import llm_models_root +from ..accuracy_core import GSM8K, LlmapiAccuracyTestHarness + +# 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", +) + +# The chat-GSM8K acceptance floors of the in-tree test on an NVFP4 KV cache (test_llm_api_pytorch.TestMiniMaxM3). +MIN_ACCEPTANCE_RATE = 0.78 +MIN_ACCEPTANCE_LENGTH = 3.3 + + +def _drain_spec_stats(llm): + drafted = accepted = steps = 0 + for s in llm.get_stats(timeout=2): + s = json.loads(s) if isinstance(s, str) else s + sd = s.get("specDecodingStats") or {} + drafted += sd.get("numDraftTokens", 0) + accepted += sd.get("numAcceptedTokens", 0) + steps += sd.get("numRequestsWithDraftTokens", 0) + return drafted, accepted, steps + + +class _MiniMaxM3Target(LlmapiAccuracyTestHarness): + MODEL_NAME = "nvidia/MiniMax-M3-NVFP4" + MODEL_PATH = f"{llm_models_root()}/MiniMax-M3-NVFP4" + EAGLE_PATH = f"{llm_models_root()}/MiniMax-M3-EAGLE3-GQA" + TP = 4 + + def _llm(self, eagle3: bool) -> 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"}, + ) + + def _gsm8k(self, eagle3: bool) -> None: + with self._llm(eagle3) as llm: + assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4 + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + if not eagle3: + return + from datasets import load_dataset + + questions = [ + r["question"] for r in load_dataset(GSM8K.DATASET_DIR, "main", split="test") + ][:200] + prompts = [ + llm.tokenizer.apply_chat_template( + [{"role": "user", "content": q}], tokenize=False, add_generation_prompt=True + ) + for q in questions + ] + _drain_spec_stats(llm) + llm.generate(prompts, SamplingParams(max_tokens=512, temperature=0)) + drafted, accepted, steps = _drain_spec_stats(llm) + assert steps > 0, "no speculative iterations recorded" + rate, length = accepted / drafted, 1 + accepted / steps + print( + f"MiniMax-M3 tp{self.TP} Eagle3 chat-GSM8K acceptance: rate={rate:.3f}, length={length:.3f}" + ) + assert rate > MIN_ACCEPTANCE_RATE, ( + f"acceptance rate {rate:.3f} <= {MIN_ACCEPTANCE_RATE}" + ) + assert length > MIN_ACCEPTANCE_LENGTH, ( + f"acceptance length {length:.3f} <= {MIN_ACCEPTANCE_LENGTH}" + ) + + +class TestModelingV2MiniMaxM3Nvfp4Sm103Tp4(_MiniMaxM3Target): + """MiniMax-M3 NVFP4 / sm_103 / tp4: four GB300s.""" + + TP = 4 + + @skip_not_sm103 + @pytest.mark.skip_less_device(4) + def test_gsm8k(self): + self._gsm8k(eagle3=False) + + @skip_not_sm103 + @pytest.mark.skip_less_device(4) + def test_eagle3_gsm8k_and_acceptance(self): + self._gsm8k(eagle3=True) + + +class TestModelingV2MiniMaxM3Nvfp4Sm103Tp8(_MiniMaxM3Target): + """MiniMax-M3 NVFP4 / sm_103 / tp8: eight GB300s over two trays (an 8-rank MPI world across the two nodes).""" + + TP = 8 + + @skip_not_sm103 + @pytest.mark.skip_less_mpi_world_size(8) + def test_gsm8k(self): + self._gsm8k(eagle3=False) + + @skip_not_sm103 + @pytest.mark.skip_less_mpi_world_size(8) + def test_eagle3_gsm8k_and_acceptance(self): + self._gsm8k(eagle3=True) diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index 1ae313f8fdad..93a6c4351dab 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -20,6 +20,8 @@ l0_gb300_multi_gpus: - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_gsm8k_identity_vs_mtp3 - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_mtp3_acceptance[modeling_v2] - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_mtp3_acceptance[stock] + - accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_gsm8k + - accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_eagle3_gsm8k_and_acceptance # Covers tests/unittest/_torch/attention/. Two sub-trees moved in here from elsewhere # under tests/unittest/_torch/ and have no entry of their own on any list, so this entry # is what picks them up: diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml index 0806efd6646a..36e1f765947e 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml @@ -17,3 +17,6 @@ l0_gb300_multi_nodes_node2_gpu8: # Kimi K3 GSM8K accuracy plus an acceptance-length gate on the TEP8 DSpark # recipe. Measured 25m45s on 2x4 GB300. - accuracy/test_kimi3.py::TestKimiK3DSpark::test_gsm8k_tep8 TIMEOUT (60) + # The MiniMax-M3 modeling_v2 tp8 target (two trays): GSM8K, and with Eagle3 GSM8K plus a chat acceptance gate. + - accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp8::test_gsm8k TIMEOUT (60) + - accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp8::test_eagle3_gsm8k_and_acceptance TIMEOUT (60) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/activation/test_modeling_v2_m3_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/activation/test_modeling_v2_m3_silu_and_mul.py new file mode 100644 index 000000000000..ea8d241d40c9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/attention/_m3_attn_runtime_tp4.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/attention/test_modeling_v2_m3_dense_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_dense_decode.py index 5a7f07733314..aab9052fa576 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_dense_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_dense_decode.py @@ -14,7 +14,7 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4 import ( # noqa: E501 - dispatch, + modeling as dispatch, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_dense_decode import ( # noqa: E501 CAPPED, diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_draft_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_draft_attention.py new file mode 100644 index 000000000000..222c0aa4fcef --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_draft_attention.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_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_103__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 + ) + + +@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 + 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_103__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_103__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..867eb5089c25 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/attention/test_modeling_v2_m3_msa_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_msa_attention.py new file mode 100644 index 000000000000..4b96f18d5c09 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..2e80fa0e980b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/attention/test_modeling_v2_m3_qkv_insert.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_qkv_insert.py new file mode 100644 index 000000000000..c570f15f772a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/attention/test_modeling_v2_m3_select_blocks.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_select_blocks.py new file mode 100644 index 000000000000..77583bdb979e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/comm/_m3_draft_ar_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_allreduce_op_matrix.py similarity index 78% rename from tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_draft_ar_norm_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_allreduce_op_matrix.py index c782ac2b782c..39e0e20b4138 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_draft_ar_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_allreduce_op_matrix.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The fast draft loop's layer boundaries (``spec_worker.draft_ar_norm``) against stock Eagle3's, bit for bit. +"""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 @@ -14,9 +16,9 @@ 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``. +``test_modeling_v2_m3_allreduce.py``. - CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_draft_ar_norm_op_matrix.py + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_allreduce_op_matrix.py """ import math @@ -188,17 +190,76 @@ def _ulps(out: torch.Tensor, ref: torch.Tensor) -> float: 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: ``spec_worker.draft_ar_norm`` with the drafter's norm module.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.spec_worker import ( # noqa: E501 - draft_ar_norm, + """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_103__tp4.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, ) - from tensorrt_llm._torch.modules.rms_norm import RMSNorm - norm = RMSNorm(hidden_size=HIDDEN, eps=EPS, dtype=torch.bfloat16).cuda() - with torch.no_grad(): - norm.weight.copy_(w) - return draft_ar_norm(ar, x, residual.clone(), norm) + 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.""" + import sys as _sys + from pathlib import Path as _Path + + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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): @@ -237,6 +298,7 @@ def _run_matrix() -> None: 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, diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_allreduce.py similarity index 55% rename from tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py rename to tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_allreduce.py index d2a966c7d5c1..4f34ae908c81 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_allreduce.py @@ -1,11 +1,11 @@ # 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. +"""Collected entry point for the ``m3_allreduce`` certification matrix. -The matrix is ``_m3_draft_ar_norm_op_matrix.py`` beside this file, its own 4-rank launcher; see ``_rank_job``. It -requires ``spec_worker.draft_ar_norm`` 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 and case. +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 @@ -21,8 +21,8 @@ # 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") +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( diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py index 8b8585dc2677..31a054642e89 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py @@ -11,10 +11,6 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.comm.m3_mc_finalize_allreduce_norm import ( # noqa: E501 - finalize_partial, -) - assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" @@ -22,17 +18,3 @@ @pytest.mark.no_xdist def test_m3_mc_allreduce_norm_op_matrix() -> None: _rank_job.run("m3_mc_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)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..4d7b4f235c59 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__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_103__tp4/embedding/test_modeling_v2_m3_embedding.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/embedding/test_modeling_v2_m3_embedding.py new file mode 100644 index 000000000000..39cedf0fff84 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/gemm/_m3_mxfp8_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/gemm/test_modeling_v2_m3_bf16_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_bf16_linear.py new file mode 100644 index 000000000000..08349a6dbe71 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/gemm/test_modeling_v2_m3_fp32_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_fp32_linear.py new file mode 100644 index 000000000000..a008ead936da --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py new file mode 100644 index 000000000000..043b4da71ff9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py new file mode 100644 index 000000000000..221cee463aa1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/moe/_m3_router.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/_m3_router.py new file mode 100644 index 000000000000..0eefbc3aa8ea --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/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_103__tp4.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_103__tp4.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_103__tp4.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_103__tp4.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_103__tp4/moe/test_modeling_v2_m3_router_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_gemv.py index baa92e7c7522..2da1ac3c6470 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_gemv.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_gemv.py @@ -1,82 +1,25 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU tests for the router entries of the MiniMax-M3 tp4 target: m3_router_gemv, m3_router_gemv_fp4 and -m3_router_tc (one kernel source, one test module).""" - -import zlib +"""GPU test for the m3_router_gemv entry of the MiniMax-M3 tp4 target.""" import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.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_103__tp4.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_103__tp4.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_103__tp4.ops.moe.m3_router_tc import ( # noqa: E501 - m3_router_tc as tc, +__extra_import_path__ = ["../../../..", "."] +from _m3_router import ( # noqa: E402, F401 + SF_SCALE, + _build, + _compare_fp4, + _compare_logits, + _seed, + gemv, + gemv_fp4, + tc, ) - -__extra_import_path__ = ["../../../.."] -from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above +from _validating import validating # noqa: E402 — needs the path declared above assert torch.cuda.is_available(), "the router entries require a CUDA device" -#: 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" - @pytest.mark.parametrize("cell", gemv.CELLS, ids=[c.why[:48] for c in gemv.CELLS]) def test_gemv_certified_cells(cell) -> None: @@ -87,62 +30,9 @@ def test_gemv_certified_cells(cell) -> None: _compare_logits(out, gemv.reference(x, w, 1, False)) -@pytest.mark.parametrize("cell", gemv_fp4.CELLS, ids=[c.why[:48] for c in gemv_fp4.CELLS]) -def test_gemv_fp4_certified_cells(cell) -> None: - spec = cell.spec - x, w = _build(spec["m"], spec["e"], spec["k"], _seed(cell)) - sf_scale = torch.tensor([spec["sf_scale"]], dtype=torch.float32, device="cuda") - with validating(gemv_fp4): - logits, q, sf = gemv_fp4(x, w, sf_scale, False) - ref_logits, _, _ = gemv_fp4.reference(x, w, sf_scale, False) - _compare_logits(logits, ref_logits) - _compare_fp4(q, sf, x, sf_scale) - - -@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_the_gemv_is_a_true_fp32_product() -> None: """And the GEMV, by contrast, is no TF32 product: the TF32 reference rejects it.""" x, w = _build(1, 128, 6144, 6) out = gemv(x, w, 1, False) with pytest.raises(AssertionError): _compare_logits(out, tc.reference(x, w, None, False)[0]) - - -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) - - -def test_guard_refuses_a_zero_global_scale() -> None: - x, w = _build(1, 128, 6144, 4) - with pytest.raises(AssertionError, match="positive and finite"): - with validating(gemv_fp4): - gemv_fp4(x, w, torch.zeros(1, dtype=torch.float32, device="cuda"), False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_gemv_fp4.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_gemv_fp4.py new file mode 100644 index 000000000000..036161f44f43 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_gemv_fp4.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_router_gemv_fp4 entry of the MiniMax-M3 tp4 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 + +assert torch.cuda.is_available(), "the router entries require a CUDA device" + + +@pytest.mark.parametrize("cell", gemv_fp4.CELLS, ids=[c.why[:48] for c in gemv_fp4.CELLS]) +def test_gemv_fp4_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], spec["e"], spec["k"], _seed(cell)) + sf_scale = torch.tensor([spec["sf_scale"]], dtype=torch.float32, device="cuda") + with validating(gemv_fp4): + logits, q, sf = gemv_fp4(x, w, sf_scale, False) + ref_logits, _, _ = gemv_fp4.reference(x, w, sf_scale, False) + _compare_logits(logits, ref_logits) + _compare_fp4(q, sf, x, sf_scale) + + +def test_guard_refuses_a_zero_global_scale() -> None: + x, w = _build(1, 128, 6144, 4) + with pytest.raises(AssertionError, match="positive and finite"): + with validating(gemv_fp4): + gemv_fp4(x, w, torch.zeros(1, dtype=torch.float32, device="cuda"), False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_tc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_tc.py new file mode 100644 index 000000000000..03ecdb80ffcd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/moe/test_modeling_v2_m3_router_tc.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_router_tc entry of the MiniMax-M3 tp4 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 + +assert torch.cuda.is_available(), "the router entries require a CUDA device" + + +@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_103__tp4/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_step.py index e6c2b67bf6c0..891889f3cd5a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_step.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_step.py @@ -313,7 +313,8 @@ def test_fragments_round_trip() -> None: 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 + # 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 @@ -643,7 +644,7 @@ def test_prepare_head_codes_the_shard_and_fails_closed(monkeypatch) -> None: 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) + 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 @@ -655,17 +656,17 @@ def corrupt(weight, *args, **kwargs): 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) + 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: - """DRAFT_HEAD_CODED = False keeps the BF16 fragments even for a shard whose coding would decode bit for bit.""" + """``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) - monkeypatch.setattr(draft_step_kernels.dispatch, "DRAFT_HEAD_CODED", False) - prepared = draft_step_kernels.prepare_head(SimpleNamespace(weight=w, bias=None), 2) + 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)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py index f66014d47aac..249679ed08fe 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py @@ -21,8 +21,7 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4 import ( # noqa: E501 - linears, - spec_worker, + modeling, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 m3_spec_tail, @@ -72,8 +71,8 @@ def _stock_accept(logits: torch.Tensor, draft: torch.Tensor): 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) + 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( @@ -268,8 +267,7 @@ def test_the_drafter_is_the_in_tree_drafter() -> None: 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 ( - spec_worker.M3Eagle3Worker._forward_draft_loop - is not Eagle3OneModelWorker._forward_draft_loop + modeling.M3Eagle3Worker._forward_draft_loop is not Eagle3OneModelWorker._forward_draft_loop ) for name in ( "_forward_linear_draft_loop", @@ -277,17 +275,15 @@ def test_the_drafter_is_the_in_tree_drafter() -> None: "sample_draft_tokens", "sample_and_accept_draft_tokens", ): - assert getattr(spec_worker.M3Eagle3Worker, name) is getattr(Eagle3OneModelWorker, name), ( - name - ) + assert getattr(modeling.M3Eagle3Worker, name) is getattr(Eagle3OneModelWorker, name), name for module, name in ( - (spec_worker, "M3DraftGemvLinearMethod"), - (spec_worker, "adopt_draft_linears"), - (spec_worker, "greedy_head_outputs"), - (linears, "M3FP8LMHeadMethod"), - (linears, "adopt_fp8_lm_heads"), + (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 = spec_worker.__name__.rsplit(".", 1)[0] + pkg = modeling.__name__.rsplit(".", 1)[0] for gone in ("kernels.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_103__tp4/spec/test_modeling_v2_m3_force_vec.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_force_vec.py new file mode 100644 index 000000000000..d8808030aba3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp4/test_modeling_v2_m3_capture_boundary.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_capture_boundary.py index 8d8835ad067e..3fc723b67501 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_capture_boundary.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_capture_boundary.py @@ -1,65 +1,103 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The decoder layer's feed-forward exit on Eagle3 capture layers (``M3DecoderLayer._feed_forward_boundary``). +"""The generic layer's feed-forward exit on Eagle3 capture layers, in both targets' bodies (``PrefillTarget._layer``, +``DecodeTarget._generic_layer``). -A capture layer with a next norm keeps the fused boundary (all-reduce + residual + next norm in one kernel) and -captures the reduced residual sum it writes; the last layer captures its raw output with the residual, as the -in-tree capture does. CPU only: the boundary and the capture are stand-ins. +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_103__tp4 import ( # noqa: E501 - layers, + modeling, ) +LAYERS = 4 -class _Layer: - """The attributes the boundary reads, with a recording capture and next-norm stand-in.""" - _fused_feed_forward_boundary = layers.M3DecoderLayer._fused_feed_forward_boundary - _feed_forward_boundary = layers.M3DecoderLayer._feed_forward_boundary +class _Spec: + def __init__(self, layers): + self.layers = set(layers) - def __init__(self, has_next_norm: bool): - self.layer_idx = 56 - self.next_layer_layernorm = object() if has_next_norm else None - self.captured = [] - self.boundaries = [] + def is_layer_capture(self, i): + return i in self.layers - def capture(self, spec_metadata, layer_id, hidden_states, residual): - self.captured.append((layer_id, hidden_states, residual)) - def _next_layer_norm(self, hidden_states, residual, fused_boundary): - self.boundaries.append(fused_boundary) - if self.next_layer_layernorm is None: - return hidden_states, residual - return ("normed", hidden_states), ("updated", hidden_states, residual) +@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) -def test_a_capture_layer_keeps_the_fused_boundary_and_captures_its_residual_sum() -> None: - layer = _Layer(has_next_norm=True) - fused = layer._fused_feed_forward_boundary(capture=True) - assert fused - h, r = torch.ones(2), torch.zeros(2) - out = layer._feed_forward_boundary(h, r, fused, True, None) - assert layer.boundaries == [True] - assert layer.captured == [(56, ("updated", h, r), None)] - assert out == (("normed", h), ("updated", h, r)) + 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: ("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: (("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 -def test_the_last_layer_captures_its_raw_output_with_the_residual() -> None: - layer = _Layer(has_next_norm=False) - fused = layer._fused_feed_forward_boundary(capture=True) - assert not fused - h, r = torch.ones(2), torch.zeros(2) - out = layer._feed_forward_boundary(h, r, fused, True, None) - assert layer.captured == [(56, h, r)] and layer.boundaries == [False] - assert out == (h, r) +@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_is_unchanged() -> None: - layer = _Layer(has_next_norm=True) - assert layer._fused_feed_forward_boundary(capture=False) - h, r = torch.ones(2), torch.zeros(2) - layer._feed_forward_boundary(h, r, True, False, None) - assert not layer.captured and layer.boundaries == [True] +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_103__tp4/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_dispatch.py index 7a06b668d5a4..1fa890f5f127 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_dispatch.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_dispatch.py @@ -1,7 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The decode target's dispatch table (``dispatch.py``): its row ranges against the catalog cells that certify -them and the kernels' own limits, the selection for every decode row count, and the miss / uncovered policy. +"""The decode target's dispatch table (``modeling.py``, imported here as ``dispatch``): its row ranges against the +catalog cells that certify them and the kernels' own limits, the selection for every decode row count, and the miss +/ uncovered policy. CPU only: nothing here launches a kernel. """ @@ -14,9 +15,7 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4 import ( # noqa: E501 - dispatch, - layers, - spec_worker, + modeling as dispatch, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 m3_attn_mega_c1, @@ -30,8 +29,8 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 mxfp8_low_m_gemm as low_m, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.modeling import ( # noqa: E501 - DecodeTarget, +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels.mx_gemv import ( # noqa: E501 + op as mx_gemv, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm import ( # noqa: E501 m3_target_head as target_head_entry, @@ -51,8 +50,6 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec import ( # noqa: E501 m3_draft_step as draft_step_entry, ) -from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.mx_gemv import op as mx_gemv -from tensorrt_llm._torch.modules.gated_mlp import GatedMLP from tensorrt_llm._torch.modules.linear import TensorParallelMode PROJECTIONS = (QKV_SPARSE, QKV, O_PROJ, SHARED_GATE_UP, SHARED_DOWN, DENSE_GATE_UP, DENSE_DOWN) @@ -292,6 +289,80 @@ def test_a_miss_logs_once_per_reason_otherwise(log) -> None: assert "m3_moe_engine at 5 rows: the engine is not built" in log.warning[0] +# ------------------------------------------------------------------ the phase split + + +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. Nothing but ``num_contexts`` is read: the metadata stand-in carries no other field.""" + 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" + + +_DECODE_FIELDS = ( + "num_tokens", + "num_seqs", + "kv_cache_manager", + "msa_idx_k_cache", + "msa_out_cache_loc", + "msa_decode_span", + "msa_block_table", + "msa_seq_lens_cuda", + "_msa_short_ok", + "msa_short_table", + "msa_proxy_max_score_view", + "msa_worst_case_max_k_tiles", + "msa_n_valid_blocks", +) + + +def _decode_metadata(num_tokens: int) -> SimpleNamespace: + """The decode projection's fields and nothing else (no ``num_contexts`` / ``num_ctx_tokens``).""" + md = SimpleNamespace(**{name: object() for name in _DECODE_FIELDS}) + md.num_tokens = num_tokens + return md + + +def test_the_decode_projection_states_the_phase_and_reads_no_count_back() -> None: + md = _decode_metadata(4) + step = dispatch.DecodeTarget(SimpleNamespace()).step_args(md) + assert step["num_contexts"] == 0 and step["num_ctx_tokens"] == 0 + assert step["metadata"] is md and step["msa_short_ok"] is md._msa_short_ok + + +@pytest.mark.parametrize( + "rows,embeds", [(49, False), (64, False), (4, True)], ids=["49_rows", "64_rows", "embeds"] +) +def test_a_wide_generation_step_runs_every_layer_generically_in_decode( + monkeypatch, rows, embeds +) -> None: + """Past the decode kernels' rows (or with embeddings) the decode target runs its own generic layer for every + layer, with its own projection, and never the decode-kernel layer.""" + layers, seen = 3, [] + core = SimpleNamespace(num_layers=layers, w={"embed": None}) + monkeypatch.setattr(dispatch, "m3_embedding", lambda ids, table: ("embedded", int(ids.numel()))) + + def generic(self, i, position_ids, hidden, residual, step, spec_metadata, capture): + seen.append((i, step["num_contexts"], step["metadata"])) + return ("hidden", i), ("residual", i) + + def kernels(self, *args): + raise AssertionError("the decode-kernel layer ran on a wide step") + + monkeypatch.setattr(dispatch.DecodeTarget, "_generic_layer", generic) + monkeypatch.setattr(dispatch.DecodeTarget, "_layer", kernels) + md = _decode_metadata(rows) + ids = torch.arange(rows, dtype=torch.int32) + out = dispatch.DecodeTarget(core).forward( + md, None if embeds else ids, None, ("embeds",) if embeds else None, None, None + ) + assert out == ("hidden", layers - 1) + assert seen == [(i, 0, md) for i in range(layers)] + + def test_uncovered_never_raises_and_logs_once(log) -> None: dispatch.set_require(True) for _ in range(3): @@ -300,14 +371,13 @@ def test_uncovered_never_raises_and_logs_once(log) -> None: class _Core: - """The parts of the text model ``DecodeTarget._embed`` reads: no embedding owners built.""" + """The parts of the core ``DecodeTarget``'s embedding and push targets read: no embedding owners built.""" def __init__(self): - self._embed_owners = {} - - @staticmethod - def embed_tokens(input_ids: torch.Tensor) -> torch.Tensor: - return input_ids.to(torch.float32)[:, None].expand(-1, 4) + self.embed_owners = {} + self.w = {"embed": torch.arange(64 * 4, dtype=torch.float32).view(64, 4)} + self.mapping = None + self.hidden = 6144 @pytest.mark.parametrize("rows", (1, 8)) @@ -316,11 +386,11 @@ def test_embedding_without_its_owner_raises_under_require(rows) -> None: with pytest.raises( dispatch.FastPathUnavailable, match=rf"'embed_firstnorm'.* {rows} rows.*owner" ): - DecodeTarget(_Core())._embed(torch.arange(rows, dtype=torch.int32), rows) + dispatch.DecodeTarget(_Core())._embed(torch.arange(rows, dtype=torch.int32), rows) def test_embedding_without_its_owner_falls_back_and_logs_once(log) -> None: - target = DecodeTarget(_Core()) + target = dispatch.DecodeTarget(_Core()) for _ in range(2): hidden, residual = target._embed(torch.arange(3, dtype=torch.int32), 3) assert residual is None and hidden.shape == (3, 4) @@ -329,7 +399,7 @@ def test_embedding_without_its_owner_falls_back_and_logs_once(log) -> None: def test_embedding_past_its_rows_is_uncovered_in_every_mode(log) -> None: dispatch.set_require(True) - hidden, residual = DecodeTarget(_Core())._embed(torch.arange(9, dtype=torch.int32), 9) + hidden, residual = dispatch.DecodeTarget(_Core())._embed(torch.arange(9, dtype=torch.int32), 9) assert residual is None and hidden.shape == (9, 4) assert len(log.info) == 1 and "embed_firstnorm" in log.info[0] and not log.warning @@ -338,41 +408,35 @@ def test_a_push_without_its_workspace_raises_under_require(monkeypatch) -> None: monkeypatch.setattr(m3_mc_allreduce, "get_workspace", lambda mapping, hidden: None) dispatch.set_require(True) with pytest.raises(dispatch.FastPathUnavailable, match=r"'attn_push'.* 4 rows"): - DecodeTarget(_Core())._attn_push_target(_Layer(), 4, 6144) + dispatch.DecodeTarget(_Core())._attn_push_target(4) def test_a_push_without_its_workspace_falls_back_and_logs_once(monkeypatch, log) -> None: monkeypatch.setattr(m3_mc_allreduce, "get_workspace", lambda mapping, hidden: None) - target = DecodeTarget(_Core()) - assert target._attn_push_target(_Layer(), 4, 6144) is None - assert target._moe_push_targets(_Layer(), 4, 6144) is None - assert target._attn_push_target(_Layer(), 4, 6144) is None + target = dispatch.DecodeTarget(_Core()) + assert target._attn_push_target(4) is None + assert target._moe_push_targets(4) is None + assert target._attn_push_target(4) is None assert len(log.warning) == 2 assert "attn_push at 4 rows" in log.warning[0] and "moe_push at 4 rows" in log.warning[1] -class _Layer: - """The parts of a decoder layer the push targets read.""" - - mapping = None - - -def test_every_dispatch_name_the_target_uses_exists() -> None: - """A fallback message or selection that names a missing ``dispatch.X`` fails only when that branch runs.""" +def test_the_target_has_no_dispatch_module_left() -> None: + """The dispatch tables live in ``modeling.py``: no module of the target still reaches a ``dispatch`` module.""" target_dir = Path(dispatch.__file__).parent - missing = [] + assert not (target_dir / "dispatch.py").exists() + stale = [] for path in sorted(target_dir.rglob("*.py")): for node in ast.walk(ast.parse(path.read_text())): + if isinstance(node, ast.ImportFrom) and any(a.name == "dispatch" for a in node.names): + stale.append(f"{path.relative_to(target_dir)}:{node.lineno}") if ( isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name) and node.value.id == "dispatch" - and not hasattr(dispatch, node.attr) ): - missing.append( - f"{path.relative_to(target_dir)}:{node.lineno}: dispatch.{node.attr}" - ) - assert not missing, missing + stale.append(f"{path.relative_to(target_dir)}:{node.lineno}: dispatch.{node.attr}") + assert not stale, stale def test_low_m_gemm_cells_follow_its_configurations() -> None: @@ -388,23 +452,8 @@ def test_low_m_gemm_cells_follow_its_configurations() -> None: def _drafter(head_weight=None): - """A drafter with only what M3Eagle3Worker._m3_drafter_supported reads: one plain Eagle3 layer with a gated MLP, - a whole-table embedding, all-reduced o_proj and down_proj and a vocab-sharded BF16 head.""" - mlp = GatedMLP.__new__(GatedMLP) - torch.nn.Module.__init__(mlp) - mlp.down_proj = SimpleNamespace(all_reduce=object()) - layer = SimpleNamespace( - _next_layer_regular=False, - mlp=mlp, - self_attn=SimpleNamespace(o_proj=SimpleNamespace(all_reduce=object())), - ) - model = SimpleNamespace( - num_layers=1, - _use_mla=False, - _eh_proj_before_attn=False, - midlayer=layer, - embed_tokens=SimpleNamespace(tp_size=1), - ) + """A drafter with only what M3Eagle3Worker._m3_drafter_supported reads: its vocab-sharded BF16 head (the + drafter's one-layer shape is validated when ``M3Eagle3Draft`` is built).""" weight = ( head_weight if head_weight is not None @@ -413,13 +462,13 @@ def _drafter(head_weight=None): lm_head = SimpleNamespace( weight=weight, bias=None, tp_mode=TensorParallelMode.COLUMN, padding_size=0, tp_size=4 ) - return SimpleNamespace(model=model, lm_head=lm_head) + return SimpleNamespace(lm_head=lm_head) def _worker(): """An M3Eagle3Worker past the step's verify: the fused verify left the first-pass gather rows and the fused first draft inputs; nothing else of the in-tree worker is read.""" - worker = spec_worker.M3Eagle3Worker.__new__(spec_worker.M3Eagle3Worker) + worker = dispatch.M3Eagle3Worker.__new__(dispatch.M3Eagle3Worker) worker.mapping = SimpleNamespace(tp_size=4, tp_rank=1, enable_attention_dp=False) worker._d2t = None worker._m3_draft_head = None @@ -441,9 +490,7 @@ def no_gpu_draft_head(monkeypatch): def test_draft_loop_without_the_prebuilt_head_raises_under_require(no_gpu_draft_head) -> None: dispatch.set_require(True) - with pytest.raises( - dispatch.FastPathUnavailable, match=r"'m3_draft_step'.* 4 rows.*prebuild_draft_head" - ): + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_draft_step'.* 4 rows.*prebuild"): _worker()._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, _drafter()) @@ -549,9 +596,7 @@ def test_an_unsupported_drafter_is_uncovered(no_gpu_draft_head, log) -> None: worker = _worker() padded = _drafter() padded.lm_head.padding_size = 16 - two_layers = _drafter() - two_layers.model.num_layers = 2 - for drafter in (padded, two_layers): + for drafter in (padded,): assert not worker.prebuild_draft_head(drafter) assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) assert log.info and not log.warning @@ -590,16 +635,14 @@ def no_gpu_target_head(monkeypatch): def _target_worker(): - worker = spec_worker.M3Eagle3Worker.__new__(spec_worker.M3Eagle3Worker) + worker = dispatch.M3Eagle3Worker.__new__(dispatch.M3Eagle3Worker) worker._m3_target_head = None return worker def test_target_head_without_the_prebuilt_copy_raises_under_require() -> None: dispatch.set_require(True) - with pytest.raises( - dispatch.FastPathUnavailable, match=r"'m3_target_head'.* 4 rows.*prebuild_target_head" - ): + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_target_head'.* 4 rows.*prebuild"): _target_worker()._m3_coded_head(_target_lm_head(), 4) @@ -688,7 +731,7 @@ def maybe_capture_hidden_states(self, layer_id, hidden_states, residual) -> None def _fc_early_worker(monkeypatch, eligible: bool = True): - worker = spec_worker.M3Eagle3Worker.__new__(spec_worker.M3Eagle3Worker) + worker = dispatch.M3Eagle3Worker.__new__(dispatch.M3Eagle3Worker) worker._m3_capture_norm = None worker._m3_fc_norm_w = None worker._m3_fc_early_drafter = None @@ -716,8 +759,8 @@ def test_the_fc_is_forked_at_the_last_captured_layer(monkeypatch, nreq) -> None: worker.capture_hidden_states(md, layer_id, None) assert not forks worker.capture_hidden_states(md, md.layers_to_capture[-1], None) - worker.capture_hidden_states(md, spec_worker.FINAL_HIDDEN_LAYER_ID, None) - assert md.captured == [1, 29, 56, spec_worker.FINAL_HIDDEN_LAYER_ID] + worker.capture_hidden_states(md, dispatch.FINAL_HIDDEN_LAYER_ID, None) + assert md.captured == [1, 29, 56, dispatch.FINAL_HIDDEN_LAYER_ID] assert forks == [(4 * nreq, drafter, None)] assert worker._m3_fc_early_stream == "stream-1" and worker._m3_fc_early_drafter is None @@ -727,9 +770,9 @@ def test_no_early_fork_with_a_final_state_capture(monkeypatch) -> None: stock place instead of forking at the last captured decoder layer.""" worker, forks = _fc_early_worker(monkeypatch) md = _CaptureMetadata(1) - md.layers_to_capture = (spec_worker.FINAL_HIDDEN_LAYER_ID, 29, 56) + md.layers_to_capture = (dispatch.FINAL_HIDDEN_LAYER_ID, 29, 56) worker.arm_capture_norm(SimpleNamespace(num_seqs=1), md, object(), object()) - for layer_id in (29, 56, spec_worker.FINAL_HIDDEN_LAYER_ID): + for layer_id in (29, 56, dispatch.FINAL_HIDDEN_LAYER_ID): worker.capture_hidden_states(md, layer_id, None) assert not forks and worker._m3_fc_early_stream is None @@ -748,89 +791,112 @@ def test_no_early_fork_past_its_rows_or_off_the_fused_verify(monkeypatch, nreq, # ------------------------------------------------------------------ the NVFP4 producer's coverage -def test_the_producer_covers_the_tp4_nvfp4_layers_only() -> None: - """``dispatch.PRODUCER`` covers a sparse layer on an NVFP4 cache at the TP4 heads; elsewhere the built-in - producers run without a miss.""" - covered = layers.producer_pw_covered - assert covered(True, 16, 1, 1) - assert not covered(False, 16, 1, 1) - assert not covered(True, 8, 1, 1) - assert not covered(True, 16, 2, 2) - assert dispatch.PRODUCER.covers("nvfp4_tp4") - - -class _FakeMsaAttn: - indexer_kv_dtype = "fp8" +def _producer_core(nvfp4: bool = True, heads: int = 16): + """What the producer dispatch reads off the core before its first precondition.""" + return SimpleNamespace(main_kv_is_nvfp4=nvfp4, heads=heads, kv_heads=1, index_heads=1) -def _producer_layer(monkeypatch, nvfp4: bool = True, heads: int = 16): - """What the producer dispatch reads before its first precondition: a fused-projection MSA sparse layer.""" - monkeypatch.setattr(layers, "MiniMaxM3MsaSparseAttention", _FakeMsaAttn) - return SimpleNamespace( - enable_fused_qkv_index_projection=True, - attn=_FakeMsaAttn(), - main_kv_is_fp8=not nvfp4, - main_kv_is_nvfp4=nvfp4, - num_heads=heads, - num_key_value_heads=1, - sparse_num_index_heads=1, - pos_embd_params=None, - ) +def test_the_producer_covers_the_tp4_nvfp4_layers_only() -> None: + """``PRODUCER`` covers a sparse layer on an NVFP4 cache at the TP4 heads; elsewhere the generic producer runs + without a miss.""" + covered = dispatch._producer_covered + assert covered(_producer_core()) + assert not covered(_producer_core(nvfp4=False)) + assert not covered(_producer_core(heads=8)) + assert not covered(SimpleNamespace(main_kv_is_nvfp4=True, heads=16, kv_heads=2, index_heads=2)) + assert dispatch.PRODUCER.covers("nvfp4_tp4") -def _produce(layer): +def _produce(core): packed = torch.zeros(4, 20 * 128, dtype=torch.bfloat16) - return layers.M3Attention._fused_fp8_qkv_indexer_norm_rope_kv_insert( - layer, packed, None, SimpleNamespace() - ) + return dispatch._sparse_qkv_producer(core, 3, packed, None, None, None, None) -def test_a_covered_producer_layer_missing_a_precondition_raises_under_require(monkeypatch) -> None: +def test_a_covered_producer_layer_missing_a_precondition_raises_under_require() -> None: """A TP4 NVFP4 sparse layer whose step lacks a producer precondition (here the position ids) is a miss.""" dispatch.set_require(True) with pytest.raises( dispatch.FastPathUnavailable, match=r"'m3_producer' covers this step at 4 rows" ): - _produce(_producer_layer(monkeypatch)) + _produce(_producer_core()) -def test_a_covered_producer_layer_without_the_fused_projection_raises_under_require( - monkeypatch, -) -> None: - layer = _producer_layer(monkeypatch) - layer.enable_fused_qkv_index_projection = False - dispatch.set_require(True) - with pytest.raises(dispatch.FastPathUnavailable, match="fuse_qkv_index_projection is disabled"): - layers.M3Attention._check_producer_projection(layer) +def _runtime(fused: bool, indexer: str = "fp8", kv: str = "fp8"): + """A model config with the runtime fields ``_validate_runtime`` reads.""" + sparse = SimpleNamespace( + algorithm="minimax_m3", + implementation="msa", + fuse_qkv_index_projection=fused, + indexer_kv_dtype=indexer, + ) + mode = SimpleNamespace( + has_fp8_kv_cache=lambda: kv == "fp8", has_fp4_kv_cache=lambda: kv == "nvfp4" + ) + return SimpleNamespace( + sparse_attention_config=sparse, quant_config=SimpleNamespace(quant_mode=mode) + ) -def test_an_uncovered_layer_without_the_fused_projection_is_not_a_miss(monkeypatch) -> None: - layer = _producer_layer(monkeypatch, nvfp4=False) - layer.enable_fused_qkv_index_projection = False +@pytest.mark.parametrize("fused", [True, False], ids=["five_way", "unfused"]) +def test_the_target_takes_both_projection_layouts(fused) -> None: + """The packed five-way projection and the separate QKV + BF16 index projections + (fuse_qkv_index_projection=False) are both this target's; a BF16 indexer or a BF16 KV cache is not.""" + dispatch._validate_runtime(_runtime(fused)) + dispatch._validate_runtime(_runtime(fused, kv="nvfp4")) + with pytest.raises(ValueError, match="indexer_kv_dtype='fp8'"): + dispatch._validate_runtime(_runtime(fused, indexer="bf16")) + with pytest.raises(ValueError, match="FP8 or NVFP4 KV cache"): + dispatch._validate_runtime(_runtime(fused, kv="bf16")) + + +@pytest.mark.parametrize("target", ["PrefillTarget", "DecodeTarget"]) +def test_an_unfused_layer_the_producer_covers_is_a_miss(target) -> None: + """As the built-in target: a layer the NVFP4 producer covers needs the five-way projection, so the unfused path is + a miss there (raised under require) before anything runs, in either target.""" dispatch.set_require(True) - layers.M3Attention._check_producer_projection(layer) + attention = getattr(dispatch, target)(_producer_core())._unfused_sparse_attention + with pytest.raises(dispatch.FastPathUnavailable, match="fuse_qkv_index_projection is disabled"): + if target == "PrefillTarget": + attention(3, None, None, None) + else: + attention(3, None, None, None, None) -def test_a_covered_producer_miss_falls_back_and_logs_once(monkeypatch, log) -> None: - layer = _producer_layer(monkeypatch) - assert _produce(layer) is None and _produce(layer) is None +def test_a_covered_producer_miss_falls_back_and_logs_once(log) -> None: + core = _producer_core() + assert _produce(core) is None and _produce(core) is None assert len(log.warning) == 1 and "m3_producer at 4 rows" in log.warning[0] @pytest.mark.parametrize("nvfp4,heads", [(False, 16), (True, 8)], ids=["fp8_cache", "tp8_heads"]) -def test_an_uncovered_producer_layer_falls_back_silently(monkeypatch, log, nvfp4, heads) -> None: +def test_an_uncovered_producer_layer_falls_back_silently(log, nvfp4, heads) -> None: dispatch.set_require(True) - assert _produce(_producer_layer(monkeypatch, nvfp4, heads)) is None + assert _produce(_producer_core(nvfp4, heads)) is None assert not log.warning def test_a_covered_producer_layer_with_the_wrong_pool_layout_is_not_supported() -> None: - """The pool-layout check the covered branch takes before the native fallback (then a miss).""" + """The pool-layout check the covered branch takes before the generic producer (then a miss).""" pool = torch.zeros(4, 2, 1, 128, 32, dtype=torch.uint8) scales = torch.zeros(4, 2, 1, 128, 8, dtype=torch.uint8) - assert not layers.producer_pw_supported(16, 1, 1, pool, scales) + assert not dispatch._producer_layout(pool, scales) good = torch.zeros(4, 2, 1, 128, 64, dtype=torch.uint8) - assert layers.producer_pw_supported(16, 1, 1, good, scales) - assert not layers.producer_pw_supported( - 16, 1, 1, good, scales.transpose(3, 4).contiguous().transpose(3, 4) - ) + assert dispatch._producer_layout(good, scales) + assert not dispatch._producer_layout(good, scales.transpose(3, 4).contiguous().transpose(3, 4)) + + +# ------------------------------------------------------------------ the shell's speculative plumbing + + +def test_the_shell_installs_guided_decoding_on_the_worker() -> None: + """With speculation the guided decoder goes to the worker (the in-tree one-engine shell's delegation); without a + worker the shell reports it cannot take one.""" + shell_cls = dispatch.ModelingV2MiniMaxM3Nvfp4Sm103Tp4 + shell = shell_cls.__new__(shell_cls) + torch.nn.Module.__init__(shell) + installed = [] + shell.spec_worker = SimpleNamespace(set_guided_decoder=lambda g: installed.append(g) or True) + decoder = object() + assert shell.set_guided_decoder(decoder) and installed == [decoder] + shell.spec_worker = None + assert shell.set_guided_decoder(decoder) is False diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_prebuild.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_prebuild.py index 34bf7bafe4c2..1a547c0039ac 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_prebuild.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_prebuild.py @@ -1,28 +1,28 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The decode kernels ``build_targets`` 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, and the index scorer with and without its early-issue schedule at every decode -query length. +"""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, and the index scorer with and without its early-issue schedule at every decode query length. 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 -from unittest import mock import pytest import torch -from torch import nn 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_103__tp4 import ( # noqa: E501 - dispatch, - layers, - linears, modeling, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 + m3_index_score as index_score_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels.mx_gemv import ( # noqa: E501 + op as mx_gemv, +) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_index_score import ( # noqa: E501 m3_index_score, ) @@ -31,13 +31,25 @@ SHARED_DOWN, SHARED_GATE_UP, ) -from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.mx_gemv import 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", + "decode_query_lens", +) @pytest.fixture(autouse=True) @@ -55,61 +67,45 @@ def __init__(self, indexer_kv_dtype: str): self.m3_config = SimpleNamespace(num_index_heads=INDEX_HEADS, sparse_index_dim=INDEX_DIM) -def _moe(layer_idx: int, shared: bool = True) -> layers.M3MoE: - """An M3MoE built by its own constructor over a stub stock block: only what prebuild_engines and the lookups - read (the builds are stubbed).""" - - def stock_init(self, model_config, aux_stream_dict, layer_idx=None): - nn.Module.__init__(self) - # One rank: the build agreement takes no collective. - self.mapping = SimpleNamespace(tp_size=1, enable_attention_dp=False) - self.shared_experts = None - - with mock.patch.object(layers.MiniMaxM3MoE, "__init__", stock_init): - moe = layers.M3MoE(None, {}, layer_idx) - moe.shared_experts = object() if shared else None - return moe - - -def _model(monkeypatch, num_layers: int = 3, indexer_kv_dtype: str = "fp8"): - """The parts of MiniMaxM3Tp4ModelingV2 prebuild_kernels reads: layers (MoE + sparse attention), the embedding - device and the speculative config.""" - monkeypatch.setattr(layers, "MiniMaxM3MsaSparseAttention", _FakeSparseAttn) - model_layers = [] - for idx in range(num_layers): - model_layers.append( - SimpleNamespace( - block_sparse_moe=_moe(idx, shared=idx != 1) if idx else None, - self_attn=SimpleNamespace( - is_sparse_attention_layer=idx != 0, attn=_FakeSparseAttn(indexer_kv_dtype) - ), - ) - ) - stub = SimpleNamespace( - layers=model_layers, - embed_tokens=SimpleNamespace(weight=torch.empty(1, device="cuda")), +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), + 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: [], ) - stub.decode_query_lens = lambda: modeling.MiniMaxM3Tp4ModelingV2.decode_query_lens(stub) - return stub - + for name in _CORE_METHODS: + setattr(core, name, getattr(modeling.MiniMaxM3Tp4ModelingV2, name).__get__(core)) + built = [] if built is None else built -def _stub_engine_builds(monkeypatch, built: list) -> None: - def build(self, num_tokens, shared=True): - built.append(("moe", self.layer_idx, num_tokens, shared)) - return ("engine", self.layer_idx, num_tokens) + 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(self, num_tokens): - built.append(("moe16", self.layer_idx, num_tokens)) - return ("engine16", self.layer_idx, 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) - monkeypatch.setattr(layers.M3MoE, "_build_engine", build) - monkeypatch.setattr(layers.M3MoE, "_build_moe16_engine", build16) + core._build_engine, core._build_moe16_engine = build, build16 + return core -def _forbid_compiles(monkeypatch) -> None: +def _forbid_compiles(monkeypatch, core=None) -> None: """Every compile / build entry point a forward could reach raises.""" import cutlass.cute as cute @@ -119,122 +115,119 @@ def boom(*args, **kwargs): monkeypatch.setattr(cute, "compile", boom) monkeypatch.setattr(mx_gemv, "ensure_compiled", boom) monkeypatch.setattr(mx_gemv, "precompile", boom) - monkeypatch.setattr(layers.M3MoE, "_build_engine", boom) - monkeypatch.setattr(layers.M3MoE, "_build_moe16_engine", boom) - monkeypatch.setattr(layers._scorer_runner(), "_compile", boom) + monkeypatch.setattr(index_score_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(monkeypatch) -> None: - model = _model(monkeypatch) - assert model.decode_query_lens() == [1, 1 + DRAFT_LEN] - model.model_config.spec_config = None - assert model.decode_query_lens() == [1] - model.model_config.spec_config = SimpleNamespace( +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 model.decode_query_lens() == [1, 3, 5] + assert core.decode_query_lens() == [1, 3, 5] -def test_prebuild_builds_every_engine_count_and_scorer_variant(monkeypatch) -> None: +def test_prebuild_builds_every_engine_count_and_scorer_variant() -> None: built = [] - _stub_engine_builds(monkeypatch, built) - monkeypatch.setattr(modeling, "adopt_m3_linears", lambda root: built.append(("linears",))) - model = _model(monkeypatch) - times = modeling.MiniMaxM3Tp4ModelingV2.prebuild_kernels(model) + core = _core(built=built) + times = core.prebuild_kernels() - assert built[0] == ("linears",) assert mx_gemv.frozen(), "prebuild must close the GEMV compiles" - # Layer 0 is dense, layer 1 has no shared expert (no engine takes it); layer 2 gets every count, in order. - assert built[1:] == [ - ("moe", 2, n, True) for n in range(1, dispatch.MOE_ENGINE_MAX_TOKENS + 1) - ] + [("moe16", 2, n) for n in range(dispatch.MOE16_MIN_TOKENS, dispatch.MOE16_MAX_TOKENS + 1)] - moe = model.layers[2].block_sparse_moe - assert set(moe._engines) == {(n, True, True) for n in dispatch.MOE_ENGINE.keys} - assert set(moe._moe16_engines) == set(dispatch.MOE16_ENGINE.keys) - - runner = layers._scorer_runner() + + # 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, layers.INDEX_SCORE_SPLIT_K, early) + (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(runner.kernel_cache) + assert want <= set(index_score_kernel.IndexScoreRunner.kernel_cache) for key in ("gemv_s", "moe_s", "scorer_s"): assert times[key] >= 0.0 - # A second build_targets builds nothing again. + # A second build builds no engine again. built.clear() - modeling.MiniMaxM3Tp4ModelingV2.prebuild_kernels(model) - assert built == [("linears",)] + core.prebuild_kernels() + assert built == [] def test_engine_lookups_after_prebuild_never_build(monkeypatch) -> None: - built = [] - _stub_engine_builds(monkeypatch, built) - moe = _moe(5) - moe.prebuild_engines() - _forbid_compiles(monkeypatch) + 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 moe.engine(n, shared=True) == ("engine", 5, n) + assert core.engine(5, n, shared=True) == ("engine", 5, n) for n in sorted(dispatch.MOE16_ENGINE.keys): - assert moe.moe16_engine(n) == ("engine16", 5, n) + assert core.moe16_engine(5, n) == ("engine16", 5, n) def test_missing_engine_is_a_miss_not_a_build(monkeypatch) -> None: - _forbid_compiles(monkeypatch) - moe = _moe(7) + core = _core() + core.moe_engines[7], core.moe16_engines[7] = {}, {} + _forbid_compiles(monkeypatch, core) dispatch.set_require(False) - assert moe.engine(3, shared=True) is None - assert moe.moe16_engine(10) is None + 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"): - moe.engine(3, shared=True) + core.engine(7, 3, shared=True) with pytest.raises(dispatch.FastPathUnavailable, match="not built before warmup"): - moe.moe16_engine(10) + core.moe16_engine(7, 10) -def test_failed_engine_build_is_logged_then_missed(monkeypatch) -> None: - def fail(self, num_tokens, shared=True): - raise ValueError("experts backend is not TRTLLM-Gen NVFP4 SwiGLU-OAI") - - monkeypatch.setattr(layers.M3MoE, "_build_engine", fail) - monkeypatch.setattr(layers.M3MoE, "_build_moe16_engine", lambda self, n: fail(self, n)) - moe = _moe(9) - moe.prebuild_engines() - assert all(v is None for v in moe._engines.values()) and len(moe._engines) == 8 +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"): - moe.engine(1, shared=True) + 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, m3_push=push, tp_size=4) + return SimpleNamespace(weight=wq.view(n, k), weight_scale=wsf, push=push) @pytest.mark.parametrize( "nk,push", - ((O_PROJ, linears.PUSH_ATTN), (SHARED_DOWN, linears.PUSH_MOE), (SHARED_GATE_UP, None)), + ((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 = linears.gemv_variants(module) + 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 == linears.PUSH_ATTN else dispatch.MOE_PUSH.keys + 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) - assert linears.prebuild_m3_gemvs(module, use_pdl) > 0 + # 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) @@ -281,12 +274,10 @@ def test_frozen_gemv_declines_an_uncompiled_kernel(monkeypatch) -> None: @pytest.mark.parametrize("indexer_kv_dtype", ("fp8", "bf16")) def test_scorer_after_prebuild_never_compiles(monkeypatch, indexer_kv_dtype) -> None: - model = _model(monkeypatch, indexer_kv_dtype=indexer_kv_dtype) - keys = layers.prebuild_index_scorer( - model.layers, model.decode_query_lens(), torch.device("cuda") - ) + 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) + _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 @@ -309,8 +300,3 @@ def test_scorer_after_prebuild_never_compiles(monkeypatch, indexer_kv_dtype) -> early, ) torch.cuda.synchronize() - - -def test_moe_keeps_its_layer_index(): - # prebuild_engines names the layer in its build and agreement messages. - assert _moe(5).layer_idx == 5 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_rank_agreement.py index 1a906d31f9d3..170597182d6b 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_rank_agreement.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_rank_agreement.py @@ -10,11 +10,9 @@ import pytest import torch -from torch import nn from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4 import ( # noqa: E501 - dispatch, - layers, + modeling, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 m3_mc_allreduce as mc, @@ -23,10 +21,10 @@ @pytest.fixture(autouse=True) def _restore_modes(): - require = dispatch.require_mode() + require = modeling.require_mode() yield - dispatch.set_require(require) - dispatch.reset_logged() + modeling.set_require(require) + modeling.reset_logged() @pytest.mark.parametrize("finalize", [False, True]) @@ -129,52 +127,45 @@ def test_tp_agree_built_two_ranks(two_rank_comm, oks): assert out == (built if all(oks) else [None, None]) -def _moe(layer_idx: int, fail: bool) -> layers.M3MoE: - """An M3MoE on two simulated TP ranks whose engine builds are stubbed (and fail on one).""" +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 stock_init(self, model_config, aux_stream_dict, layer_idx=None): - nn.Module.__init__(self) - self.mapping = _mapping(2) - self.shared_experts = None - - with mock.patch.object(layers.MiniMaxM3MoE, "__init__", stock_init): - moe = layers.M3MoE(None, {}, layer_idx) - moe.shared_experts = object() - - def build(num_tokens, shared=True): + def build(i, num_tokens, shared=True): if fail: raise RuntimeError("compile failed on this rank") return ("engine", num_tokens) - moe._build_engine = build - moe._build_moe16_engine = lambda num_tokens: build(num_tokens) - return moe + 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): - moes = [_moe(3, f) for f in fail] - _run_ranks([m.prebuild_engines for m in moes]) - keys = sorted(moes[0]._engines) + sorted(moes[0]._moe16_engines) + 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 ("_engines", "_moe16_engines"): - a, b = getattr(moes[0], engines), getattr(moes[1], engines) + 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 built-in MoE. - dispatch.set_require(False) - for n in sorted(dispatch.MOE_ENGINE.keys): - got = [m.engine(n, shared=True) for m in moes] + # 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): - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.spec_worker import ( # noqa: E501 - M3Eagle3Worker, - ) - + 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) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/activation/test_modeling_v2_m3_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/activation/test_modeling_v2_m3_silu_and_mul.py new file mode 100644 index 000000000000..470ed6d7cbd1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/attention/_m3_attn_runtime_tp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/attention/test_modeling_v2_m3_dense_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_dense_decode.py index d7d912e26bfd..9075ab0941f7 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_dense_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_dense_decode.py @@ -14,7 +14,7 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 - dispatch, + modeling as dispatch, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_dense_decode import ( # noqa: E501 CAPPED, diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_draft_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_draft_attention.py new file mode 100644 index 000000000000..bfa051f8c35b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_draft_attention.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_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_103__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 = { + (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 + ) + + +@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 + 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_103__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_103__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..241f07fd7136 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/attention/test_modeling_v2_m3_fused_score_select.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_fused_score_select.py index 9b1edb6f77dd..f8e2bf4bfacf 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_fused_score_select.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_fused_score_select.py @@ -16,7 +16,7 @@ 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_103__tp8 import ( # noqa: E501 - dispatch, + modeling as dispatch, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 m3_fused_score_select as fss, diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_msa_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_msa_attention.py new file mode 100644 index 000000000000..060daf650f9e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/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_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_103__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"): ((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_103__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..9d16b5bb945a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/attention/test_modeling_v2_m3_qkv_insert.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_qkv_insert.py new file mode 100644 index 000000000000..eecd77e17f8d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/attention/test_modeling_v2_m3_select_blocks.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_select_blocks.py new file mode 100644 index 000000000000..c392993acf4a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_sparse_decode.py index 3551a46ca09c..3e6b37f90d8d 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_sparse_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_sparse_decode.py @@ -123,10 +123,10 @@ def test_guard_refuses_a_false_identity_promise() -> None: def test_a_selected_step_the_lc8_kernel_declines_is_a_miss() -> None: - """A selected table at the TP8 heads is the one-launch kernel's (dispatch.SPARSE_DECODE_LC8): an output it cannot - write (here a strided view) raises under require instead of silently taking the split-K chain.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( - dispatch, + """A selected table at the TP8 heads is the one-launch kernel's (modeling's SPARSE_DECODE_LC8): an output it + cannot write (here a strided view) raises under require instead of silently taking the split-K chain.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 + modeling as dispatch, ) case = build_step((9000, 40000), 4, identity=False, seed=31) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_draft_ar_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_allreduce_op_matrix.py similarity index 79% rename from tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_draft_ar_norm_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_allreduce_op_matrix.py index 271bc5d6fd3a..7b8a360875cd 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_draft_ar_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_allreduce_op_matrix.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The fast draft loop's layer boundaries (``spec_worker.draft_ar_norm``) against stock Eagle3's, bit for bit. +"""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 @@ -17,9 +19,9 @@ 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_draft_ar_norm.py``. +collected entry point is ``test_modeling_v2_m3_allreduce.py``. - CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_draft_ar_norm_op_matrix.py + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_allreduce_op_matrix.py """ import math @@ -202,17 +204,76 @@ def _ulps(out: torch.Tensor, ref: torch.Tensor) -> float: 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: ``spec_worker.draft_ar_norm`` with the drafter's norm module.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.spec_worker import ( # noqa: E501 - draft_ar_norm, + """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_103__tp8.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, ) - from tensorrt_llm._torch.modules.rms_norm import RMSNorm - norm = RMSNorm(hidden_size=HIDDEN, eps=EPS, dtype=torch.bfloat16).cuda() - with torch.no_grad(): - norm.weight.copy_(w) - return draft_ar_norm(ar, x, residual.clone(), norm) + 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.""" + import sys as _sys + from pathlib import Path as _Path + + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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): @@ -251,6 +312,7 @@ def _run_matrix() -> None: 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, diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_draft_ar_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_allreduce.py similarity index 66% rename from tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_draft_ar_norm.py rename to tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_allreduce.py index d0d521428d12..413bc0eec59b 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_draft_ar_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_allreduce.py @@ -1,11 +1,11 @@ # 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. +"""Collected entry point for the ``m3_allreduce`` certification matrix. -The matrix is ``_m3_draft_ar_norm_op_matrix.py`` beside this file, its own launcher; see ``_rank_job``. It requires -``spec_worker.draft_ar_norm`` to equal stock Eagle3's boundary (plain all-reduce, then the ``RMSNorm`` module with the -residual) 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 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 @@ -27,8 +27,8 @@ # 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_draft_ar_norm_is_stock_eagle3(world_size) -> None: - out = _rank_job.run("m3_draft_ar_norm", world_size) +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( diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py index 31ebc26c4a6d..8bae5f075307 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py @@ -17,10 +17,6 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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" @@ -44,17 +40,3 @@ def test_the_cells_are_world_independent() -> None: ): 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_103__tp8/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..2a8a2672b3de --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/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_103__tp8.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_103__tp8.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_103__tp8/embedding/test_modeling_v2_m3_embedding.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/embedding/test_modeling_v2_m3_embedding.py new file mode 100644 index 000000000000..08e4f32a7761 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/gemm/_m3_mxfp8_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/gemm/test_modeling_v2_m3_bf16_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_bf16_linear.py new file mode 100644 index 000000000000..254bd62efb7f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/gemm/test_modeling_v2_m3_draft_lowm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_draft_lowm.py index a8eb5a5c1425..7fb3cad9c0e5 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_draft_lowm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_draft_lowm.py @@ -3,22 +3,19 @@ """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. -Every cell is held to the fp64 reference within the derived band. The target's linear method -(``kernels/m3_draft_lowm.DraftLowMLinearMethod``) is held on a real ``Linear`` (BF16, no bias): a covered row count is -the entry's call bit for bit, an uncovered one the in-tree method's. +Every cell is held to the fp64 reference within the derived band. The target's drafter linear (``_draft_linear`` in +modeling.py) is held on real BF16 weights: a covered row count is the entry's call bit for bit, an uncovered one the +in-tree BF16 GEMM's (``m3_bf16_linear``). """ import zlib +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_103__tp8 import ( # noqa: E501 - dispatch, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 - m3_draft_lowm, + modeling as dispatch, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.gemm.m3_draft_lowm import ( # noqa: E501 TACTICS, @@ -28,7 +25,6 @@ m3_draft_lowm as op, ) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl -from tensorrt_llm._torch.modules.linear import Linear, UnquantizedLinearMethod __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 @@ -103,79 +99,45 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res assert torch.equal(plain.view(torch.int16), out.view(torch.int16)) -def _linear(nk, seed: int, bias: bool = False) -> Linear: - n, k = nk - linear = Linear(k, n, bias=bias, dtype=torch.bfloat16).cuda() - _, w = build(n, k, 1, seed) - linear.weight.data.copy_(w) - if bias: - linear.bias.data.zero_() - return linear - - -class _Drafter(nn.Module): - """A module tree holding drafter linears, as ``install`` / ``prebuild`` walk it.""" - - def __init__(self, *linears: Linear): - super().__init__() - self.linears = nn.ModuleList(linears) +def _drafter(weight: torch.Tensor): + """A drafter whose every linear is ``weight`` (``_prebuild_draft_lowm`` compiles per distinct shape).""" + return SimpleNamespace(w={key: weight for key in ("qkv", "o", "gu", "dn", "fc")}) @pytest.mark.parametrize("nk", sorted(DRAFTER_LINEARS), ids=lambda nk: DRAFTER_LINEARS[nk]) -def test_the_linear_method_is_the_entry_and_falls_back_past_its_rows(nk) -> None: +def test_the_drafter_linear_is_the_entry_and_falls_back_past_its_rows(nk) -> None: n, k = nk - linear = _linear(nk, seed=n + k) - drafter = _Drafter(linear) - assert m3_draft_lowm.install(drafter) == 1 - assert isinstance(linear.quant_method, m3_draft_lowm.DraftLowMLinearMethod) - m3_draft_lowm.prebuild(drafter) + _, weight = build(n, k, 1, seed=n + k) + dispatch._prebuild_draft_lowm(_drafter(weight)) use_pdl = get_env_enable_pdl() last = TACTICS[nk][-1][0] dispatch.set_require(True) for m in range(1, last + 1): assert dispatch.DRAFT_LOWM.covers((nk, m)) x, _ = build(n, k, m, seed=m) - got = linear.quant_method.apply(linear, x, None) + got = dispatch._draft_linear(x, weight) kind, tactic = tactic_for(n, k, m) - want = op( - x, linear.weight, torch.empty(m, n, dtype=x.dtype, device="cuda"), kind, tactic, use_pdl - ) + want = op(x, weight, 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 - assert torch.equal(linear(x).view(torch.int16), want.view(torch.int16)), m - # Past the table: the in-tree method, silently, even under require. + # Past the table: the in-tree BF16 GEMM, even under require. for m in (last + 1, 2 * last + 3): assert not dispatch.DRAFT_LOWM.covers((nk, m)) x, _ = build(n, k, m, seed=m) - got = linear.quant_method.apply(linear, x, None) - stock = UnquantizedLinearMethod().apply(linear, x, None) + got = dispatch._draft_linear(x, weight) + stock = dispatch.m3_bf16_linear(x, weight) assert torch.equal(got.view(torch.int16), stock.view(torch.int16)), m -def test_install_takes_only_covered_bias_free_bf16_linears() -> None: - covered = _linear((6144, 1024), seed=1) - biased = _linear((6144, 1024), seed=2, bias=True) - other = _linear((256, 512), seed=3) - drafter = _Drafter(covered, biased, other) - assert m3_draft_lowm.install(drafter) == 1 - assert isinstance(covered.quant_method, m3_draft_lowm.DraftLowMLinearMethod) - assert type(biased.quant_method) is UnquantizedLinearMethod - assert type(other.quant_method) is UnquantizedLinearMethod - # A second install finds nothing left to take. - assert m3_draft_lowm.install(drafter) == 0 - - def test_a_covered_call_not_prebuilt_raises_under_require(monkeypatch) -> None: - monkeypatch.setattr(m3_draft_lowm, "_BUILT", set()) - linear = _linear((6144, 1024), seed=4) - m3_draft_lowm.install(_Drafter(linear)) - x, _ = build(6144, 1024, 1, seed=5) + monkeypatch.setattr(dispatch, "_DRAFT_LOWM_BUILT", set()) + x, weight = build(6144, 1024, 1, seed=5) dispatch.set_require(True) with pytest.raises(dispatch.FastPathUnavailable, match="not compiled before warmup"): - linear.quant_method.apply(linear, x, None) - # Outside require it falls back to the in-tree method. + dispatch._draft_linear(x, weight) + # Outside require it falls back to the in-tree BF16 GEMM. dispatch.set_require(False) - got = linear.quant_method.apply(linear, x, None) - stock = UnquantizedLinearMethod().apply(linear, x, None) + got = dispatch._draft_linear(x, weight) + stock = dispatch.m3_bf16_linear(x, weight) assert torch.equal(got.view(torch.int16), stock.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_fp32_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_fp32_linear.py new file mode 100644 index 000000000000..4daa50048d12 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py new file mode 100644 index 000000000000..9ba7fa37d475 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py new file mode 100644 index 000000000000..dd7ede34f0b8 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/moe/_m3_router.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/_m3_router.py new file mode 100644 index 000000000000..f5d91aef73a7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__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_103__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_103__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_103__tp8/moe/test_modeling_v2_m3_router_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_gemv.py index e85564fc0fb4..2e923b4ec114 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_gemv.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_gemv.py @@ -1,82 +1,25 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU tests for the router entries of the MiniMax-M3 tp8 target: m3_router_gemv, m3_router_gemv_fp4 and -m3_router_tc (one kernel source, one test module).""" - -import zlib +"""GPU test for the m3_router_gemv entry of the MiniMax-M3 tp8 target.""" import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__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_103__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_103__tp8.ops.moe.m3_router_tc import ( # noqa: E501 - m3_router_tc as tc, +__extra_import_path__ = ["../../../..", "."] +from _m3_router import ( # noqa: E402, F401 + SF_SCALE, + _build, + _compare_fp4, + _compare_logits, + _seed, + gemv, + gemv_fp4, + tc, ) - -__extra_import_path__ = ["../../../.."] -from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above +from _validating import validating # noqa: E402 — needs the path declared above assert torch.cuda.is_available(), "the router entries require a CUDA device" -#: 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" - @pytest.mark.parametrize("cell", gemv.CELLS, ids=[c.why[:48] for c in gemv.CELLS]) def test_gemv_certified_cells(cell) -> None: @@ -87,62 +30,9 @@ def test_gemv_certified_cells(cell) -> None: _compare_logits(out, gemv.reference(x, w, 1, False)) -@pytest.mark.parametrize("cell", gemv_fp4.CELLS, ids=[c.why[:48] for c in gemv_fp4.CELLS]) -def test_gemv_fp4_certified_cells(cell) -> None: - spec = cell.spec - x, w = _build(spec["m"], spec["e"], spec["k"], _seed(cell)) - sf_scale = torch.tensor([spec["sf_scale"]], dtype=torch.float32, device="cuda") - with validating(gemv_fp4): - logits, q, sf = gemv_fp4(x, w, sf_scale, False) - ref_logits, _, _ = gemv_fp4.reference(x, w, sf_scale, False) - _compare_logits(logits, ref_logits) - _compare_fp4(q, sf, x, sf_scale) - - -@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_the_gemv_is_a_true_fp32_product() -> None: """And the GEMV, by contrast, is no TF32 product: the TF32 reference rejects it.""" x, w = _build(1, 128, 6144, 6) out = gemv(x, w, 1, False) with pytest.raises(AssertionError): _compare_logits(out, tc.reference(x, w, None, False)[0]) - - -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) - - -def test_guard_refuses_a_zero_global_scale() -> None: - x, w = _build(1, 128, 6144, 4) - with pytest.raises(AssertionError, match="positive and finite"): - with validating(gemv_fp4): - gemv_fp4(x, w, torch.zeros(1, dtype=torch.float32, device="cuda"), False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_gemv_fp4.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_gemv_fp4.py new file mode 100644 index 000000000000..11e9982c8dcc --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_gemv_fp4.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_router_gemv_fp4 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 + +assert torch.cuda.is_available(), "the router entries require a CUDA device" + + +@pytest.mark.parametrize("cell", gemv_fp4.CELLS, ids=[c.why[:48] for c in gemv_fp4.CELLS]) +def test_gemv_fp4_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], spec["e"], spec["k"], _seed(cell)) + sf_scale = torch.tensor([spec["sf_scale"]], dtype=torch.float32, device="cuda") + with validating(gemv_fp4): + logits, q, sf = gemv_fp4(x, w, sf_scale, False) + ref_logits, _, _ = gemv_fp4.reference(x, w, sf_scale, False) + _compare_logits(logits, ref_logits) + _compare_fp4(q, sf, x, sf_scale) + + +def test_guard_refuses_a_zero_global_scale() -> None: + x, w = _build(1, 128, 6144, 4) + with pytest.raises(AssertionError, match="positive and finite"): + with validating(gemv_fp4): + gemv_fp4(x, w, torch.zeros(1, dtype=torch.float32, device="cuda"), False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_tc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_tc.py new file mode 100644 index 000000000000..0b3638ee02eb --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/moe/test_modeling_v2_m3_router_tc.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_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 + +assert torch.cuda.is_available(), "the router entries require a CUDA device" + + +@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_103__tp8/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_step.py index 3e76283f4d90..339b61ce3dc0 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_step.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_step.py @@ -584,7 +584,11 @@ def test_prepare_head_streams_the_bf16_fragments_at_every_k() -> None: assert torch.equal(head_weight(prepared.wf, k).view(torch.int16), w.view(torch.int16)) with pytest.raises(NotImplementedError): head_weight(prepared.wf, 96, torch.zeros(2, dtype=torch.int32, device="cuda")) - assert not hasattr(draft_step_kernels.dispatch, "DRAFT_HEAD_CODED") + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 + modeling, + ) + + assert not hasattr(modeling, "DRAFT_HEAD_CODED") assert all(not c.spec["coded"] for c in op.CELLS) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_exact_numerics.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_exact_numerics.py index 3916557e39a5..f705462595cc 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_exact_numerics.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_exact_numerics.py @@ -5,13 +5,13 @@ * 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 BF16 linears on fixed-tactic BF16 low-M - kernels at the rows ``dispatch.DRAFT_LOWM`` covers, held to the fp64 product and to the entry by - test_modeling_v2_m3_draft_lowm.py); 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). +* Drafter: the worker runs the in-tree drafter's arithmetic and weights (its BF16 linears on fixed-tactic BF16 low-M + kernels at the rows ``DRAFT_LOWM`` covers, held to the fp64 product and to the entry by + test_modeling_v2_m3_draft_lowm.py); 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 (8 x 25008 shards), 3 captured layers, draft length 3. """ @@ -23,11 +23,9 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 - linears, - spec_worker, + modeling, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 - m3_draft_lowm, m3_spec_tail, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 @@ -39,7 +37,7 @@ ) 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, UnquantizedLinearMethod +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 @@ -75,8 +73,8 @@ def _stock_accept(logits: torch.Tensor, draft: torch.Tensor): 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) + 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( @@ -269,11 +267,9 @@ def test_draft_tokens_through_the_full_vocab_head(rows) -> None: 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, the drafter linears' low-M method is the in-tree BF16 method's subclass, and the fused head - streams the BF16 head weight itself (no coded copy).""" + left in the target, and the fused head streams the BF16 head weight itself (no coded copy).""" assert ( - spec_worker.M3Eagle3Worker._forward_draft_loop - is not Eagle3OneModelWorker._forward_draft_loop + modeling.M3Eagle3Worker._forward_draft_loop is not Eagle3OneModelWorker._forward_draft_loop ) for name in ( "_forward_linear_draft_loop", @@ -281,25 +277,23 @@ def test_the_drafter_is_the_in_tree_drafter() -> None: "sample_draft_tokens", "sample_and_accept_draft_tokens", ): - assert getattr(spec_worker.M3Eagle3Worker, name) is getattr(Eagle3OneModelWorker, name), ( - name - ) + assert getattr(modeling.M3Eagle3Worker, name) is getattr(Eagle3OneModelWorker, name), name for module, name in ( - (spec_worker, "M3DraftGemvLinearMethod"), - (spec_worker, "adopt_draft_linears"), - (spec_worker, "greedy_head_outputs"), - (linears, "M3FP8LMHeadMethod"), - (linears, "adopt_fp8_lm_heads"), + (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}" - assert issubclass(m3_draft_lowm.DraftLowMLinearMethod, UnquantizedLinearMethod) - assert not hasattr(spec_worker.M3Eagle3Worker, "prebuild_target_head") - pkg = spec_worker.__name__.rsplit(".", 1)[0] + assert not hasattr(modeling.M3Eagle3Worker, "prebuild_target_head") + pkg = modeling.__name__.rsplit(".", 1)[0] for gone in ( "kernels.m3_draft_head", "ops.spec.m3_draft_head", "ops.gemm.m3_draft_gemv", "kernels.m3_target_head", "ops.gemm.m3_target_head", + "kernels.m3_draft_lowm", ): 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_103__tp8/spec/test_modeling_v2_m3_force_vec.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_force_vec.py new file mode 100644 index 000000000000..fe0d05f52d1a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__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_103__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_103__tp8/test_modeling_v2_m3_capture_boundary.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_capture_boundary.py index ce69ccb0a568..455c283718ad 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_capture_boundary.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_capture_boundary.py @@ -1,65 +1,103 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The decoder layer's feed-forward exit on Eagle3 capture layers (``M3DecoderLayer._feed_forward_boundary``). +"""The generic layer's feed-forward exit on Eagle3 capture layers, in both targets' bodies (``PrefillTarget._layer``, +``DecodeTarget._generic_layer``). -A capture layer with a next norm keeps the fused boundary (all-reduce + residual + next norm in one kernel) and -captures the reduced residual sum it writes; the last layer captures its raw output with the residual, as the -in-tree capture does. CPU only: the boundary and the capture are stand-ins. +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_103__tp8 import ( # noqa: E501 - layers, + modeling, ) +LAYERS = 4 -class _Layer: - """The attributes the boundary reads, with a recording capture and next-norm stand-in.""" - _fused_feed_forward_boundary = layers.M3DecoderLayer._fused_feed_forward_boundary - _feed_forward_boundary = layers.M3DecoderLayer._feed_forward_boundary +class _Spec: + def __init__(self, layers): + self.layers = set(layers) - def __init__(self, has_next_norm: bool): - self.layer_idx = 56 - self.next_layer_layernorm = object() if has_next_norm else None - self.captured = [] - self.boundaries = [] + def is_layer_capture(self, i): + return i in self.layers - def capture(self, spec_metadata, layer_id, hidden_states, residual): - self.captured.append((layer_id, hidden_states, residual)) - def _next_layer_norm(self, hidden_states, residual, fused_boundary): - self.boundaries.append(fused_boundary) - if self.next_layer_layernorm is None: - return hidden_states, residual - return ("normed", hidden_states), ("updated", hidden_states, residual) +@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) -def test_a_capture_layer_keeps_the_fused_boundary_and_captures_its_residual_sum() -> None: - layer = _Layer(has_next_norm=True) - fused = layer._fused_feed_forward_boundary(capture=True) - assert fused - h, r = torch.ones(2), torch.zeros(2) - out = layer._feed_forward_boundary(h, r, fused, True, None) - assert layer.boundaries == [True] - assert layer.captured == [(56, ("updated", h, r), None)] - assert out == (("normed", h), ("updated", h, r)) + 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: ("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, push=None: (("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 -def test_the_last_layer_captures_its_raw_output_with_the_residual() -> None: - layer = _Layer(has_next_norm=False) - fused = layer._fused_feed_forward_boundary(capture=True) - assert not fused - h, r = torch.ones(2), torch.zeros(2) - out = layer._feed_forward_boundary(h, r, fused, True, None) - assert layer.captured == [(56, h, r)] and layer.boundaries == [False] - assert out == (h, r) +@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_is_unchanged() -> None: - layer = _Layer(has_next_norm=True) - assert layer._fused_feed_forward_boundary(capture=False) - h, r = torch.ones(2), torch.zeros(2) - layer._feed_forward_boundary(h, r, True, False, None) - assert not layer.captured and layer.boundaries == [True] +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_103__tp8/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_dispatch.py index 4968e98e0177..2c098e15ee5b 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_dispatch.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_dispatch.py @@ -1,7 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The tp8 decode target's dispatch table (``dispatch.py``): its row ranges against the catalog cells that certify -them and the kernels' own limits, the selection for every decode row count, and the miss / uncovered policy. +"""The tp8 decode target's dispatch table (``modeling.py``, imported here as ``dispatch``): its row ranges against the +catalog cells that certify them and the kernels' own limits, the selection for every decode row count, and the miss / +uncovered policy. CPU only: nothing here launches a kernel. """ @@ -15,14 +16,10 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 - dispatch, - layers, - linears, - spec_worker, + modeling as dispatch, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 m3_attn_lc8, - m3_draft_lowm, m3_draft_step, m3_fused_score_select, m3_mc_allreduce, @@ -32,9 +29,6 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels.mx_gemv import ( # noqa: E501 op as mx_gemv, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.modeling import ( # noqa: E501 - DecodeTarget, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention import ( # noqa: E501 m3_fused_score_select as fused_score_select_entry, ) @@ -54,7 +48,6 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec import ( # noqa: E501 m3_draft_step as draft_step_entry, ) -from tensorrt_llm._torch.modules.gated_mlp import GatedMLP from tensorrt_llm._torch.modules.linear import TensorParallelMode _TARGET = ( @@ -158,9 +151,9 @@ def test_the_removed_tp4_paths_are_gone() -> None: ): assert importlib.util.find_spec(f"{_TARGET}.{module}") is None, module for method in ("prebuild_target_head", "_m3_coded_head"): - assert not hasattr(spec_worker.M3Eagle3Worker, method), method + assert not hasattr(dispatch.M3Eagle3Worker, method), method for method in ("moe16_engine", "_build_moe16_engine"): - assert not hasattr(layers.M3MoE, method), method + assert not hasattr(dispatch.MiniMaxM3Tp8ModelingV2, method), method # ------------------------------------------------------------------ the ranges against the kernels' own limits @@ -187,8 +180,7 @@ def test_the_shared_expert_projections_take_no_gemv(nk, push) -> None: """At TP8 the shared expert (768 x 6144 gate_up, 6144 x 384 down) runs inside the MoE engine on decode rows and on the generic GEMM otherwise: the GEMV has no config for it and the dispatch never routes a call there.""" assert mx_gemv.get_config(*nk) is None - module = SimpleNamespace(weight=torch.zeros(1, 1).expand(*nk), m3_push=push, tp_size=8) - assert linears.gemv_variants(module) == [] + assert dispatch._gemv_variants(*nk, push, 8) == [] def test_the_plans_are_the_all_reduce_buckets() -> None: @@ -405,6 +397,80 @@ def test_a_miss_logs_once_per_reason_otherwise(log) -> None: assert "m3_moe_engine at 5 rows: the engine is not built" in log.warning[0] +# ------------------------------------------------------------------ the phase split + + +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. Nothing but ``num_contexts`` is read: the metadata stand-in carries no other field.""" + core = SimpleNamespace(_prefill="prefill", _decode="decode") + select = dispatch.MiniMaxM3Tp8ModelingV2._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" + + +_DECODE_FIELDS = ( + "num_tokens", + "num_seqs", + "kv_cache_manager", + "msa_idx_k_cache", + "msa_out_cache_loc", + "msa_decode_span", + "msa_block_table", + "msa_seq_lens_cuda", + "_msa_short_ok", + "msa_short_table", + "msa_proxy_max_score_view", + "msa_worst_case_max_k_tiles", + "msa_n_valid_blocks", +) + + +def _decode_metadata(num_tokens: int) -> SimpleNamespace: + """The decode projection's fields and nothing else (no ``num_contexts`` / ``num_ctx_tokens``).""" + md = SimpleNamespace(**{name: object() for name in _DECODE_FIELDS}) + md.num_tokens = num_tokens + return md + + +def test_the_decode_projection_states_the_phase_and_reads_no_count_back() -> None: + md = _decode_metadata(4) + step = dispatch.DecodeTarget(SimpleNamespace()).step_args(md) + assert step["num_contexts"] == 0 and step["num_ctx_tokens"] == 0 + assert step["metadata"] is md and step["msa_short_ok"] is md._msa_short_ok + + +@pytest.mark.parametrize( + "rows,embeds", [(49, False), (64, False), (4, True)], ids=["49_rows", "64_rows", "embeds"] +) +def test_a_wide_generation_step_runs_every_layer_generically_in_decode( + monkeypatch, rows, embeds +) -> None: + """Past the decode kernels' rows (or with embeddings) the decode target runs its own generic layer for every + layer, with its own projection, and never the decode-kernel layer.""" + layers, seen = 3, [] + core = SimpleNamespace(num_layers=layers, w={"embed": None}) + monkeypatch.setattr(dispatch, "m3_embedding", lambda ids, table: ("embedded", int(ids.numel()))) + + def generic(self, i, position_ids, hidden, residual, step, spec_metadata, capture): + seen.append((i, step["num_contexts"], step["metadata"])) + return ("hidden", i), ("residual", i) + + def kernels(self, *args): + raise AssertionError("the decode-kernel layer ran on a wide step") + + monkeypatch.setattr(dispatch.DecodeTarget, "_generic_layer", generic) + monkeypatch.setattr(dispatch.DecodeTarget, "_layer", kernels) + md = _decode_metadata(rows) + ids = torch.arange(rows, dtype=torch.int32) + out = dispatch.DecodeTarget(core).forward( + md, None if embeds else ids, None, ("embeds",) if embeds else None, None, None + ) + assert out == ("hidden", layers - 1) + assert seen == [(i, 0, md) for i in range(layers)] + + def test_uncovered_never_raises_and_logs_once(log) -> None: dispatch.set_require(True) for _ in range(3): @@ -413,14 +479,13 @@ def test_uncovered_never_raises_and_logs_once(log) -> None: class _Core: - """The parts of the text model ``DecodeTarget._embed`` reads: no embedding owners built.""" + """The parts of the core ``DecodeTarget``'s embedding and push targets read: no embedding owners built.""" def __init__(self): - self._embed_owners = {} - - @staticmethod - def embed_tokens(input_ids: torch.Tensor) -> torch.Tensor: - return input_ids.to(torch.float32)[:, None].expand(-1, 4) + self.embed_owners = {} + self.w = {"embed": torch.arange(64 * 4, dtype=torch.float32).view(64, 4)} + self.mapping = None + self.hidden = 6144 @pytest.mark.parametrize("rows", (1, 8)) @@ -429,11 +494,11 @@ def test_embedding_without_its_owner_raises_under_require(rows) -> None: with pytest.raises( dispatch.FastPathUnavailable, match=rf"'embed_firstnorm'.* {rows} rows.*owner" ): - DecodeTarget(_Core())._embed(torch.arange(rows, dtype=torch.int32), rows) + dispatch.DecodeTarget(_Core())._embed(torch.arange(rows, dtype=torch.int32), rows) def test_embedding_without_its_owner_falls_back_and_logs_once(log) -> None: - target = DecodeTarget(_Core()) + target = dispatch.DecodeTarget(_Core()) for _ in range(2): hidden, residual = target._embed(torch.arange(3, dtype=torch.int32), 3) assert residual is None and hidden.shape == (3, 4) @@ -442,7 +507,7 @@ def test_embedding_without_its_owner_falls_back_and_logs_once(log) -> None: def test_embedding_past_its_rows_is_uncovered_in_every_mode(log) -> None: dispatch.set_require(True) - hidden, residual = DecodeTarget(_Core())._embed(torch.arange(9, dtype=torch.int32), 9) + hidden, residual = dispatch.DecodeTarget(_Core())._embed(torch.arange(9, dtype=torch.int32), 9) assert residual is None and hidden.shape == (9, 4) assert len(log.info) == 1 and "embed_firstnorm" in log.info[0] and not log.warning @@ -451,41 +516,35 @@ def test_a_push_without_its_workspace_raises_under_require(monkeypatch) -> None: monkeypatch.setattr(m3_mc_allreduce, "get_workspace", lambda mapping, hidden: None) dispatch.set_require(True) with pytest.raises(dispatch.FastPathUnavailable, match=r"'attn_push'.* 4 rows"): - DecodeTarget(_Core())._attn_push_target(_Layer(), 4, 6144) + dispatch.DecodeTarget(_Core())._attn_push_target(4) def test_a_push_without_its_workspace_falls_back_and_logs_once(monkeypatch, log) -> None: monkeypatch.setattr(m3_mc_allreduce, "get_workspace", lambda mapping, hidden: None) - target = DecodeTarget(_Core()) - assert target._attn_push_target(_Layer(), 4, 6144) is None - assert target._moe_push_targets(_Layer(), 4, 6144) is None - assert target._attn_push_target(_Layer(), 4, 6144) is None + target = dispatch.DecodeTarget(_Core()) + assert target._attn_push_target(4) is None + assert target._moe_push_targets(4) is None + assert target._attn_push_target(4) is None assert len(log.warning) == 2 assert "attn_push at 4 rows" in log.warning[0] and "moe_push at 4 rows" in log.warning[1] -class _Layer: - """The parts of a decoder layer the push targets read.""" - - mapping = None - - -def test_every_dispatch_name_the_target_uses_exists() -> None: - """A fallback message or selection that names a missing ``dispatch.X`` fails only when that branch runs.""" +def test_the_target_has_no_dispatch_module_left() -> None: + """The dispatch tables live in ``modeling.py``: no module of the target still reaches a ``dispatch`` module.""" target_dir = Path(dispatch.__file__).parent - missing = [] + assert not (target_dir / "dispatch.py").exists() + stale = [] for path in sorted(target_dir.rglob("*.py")): for node in ast.walk(ast.parse(path.read_text())): + if isinstance(node, ast.ImportFrom) and any(a.name == "dispatch" for a in node.names): + stale.append(f"{path.relative_to(target_dir)}:{node.lineno}") if ( isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name) and node.value.id == "dispatch" - and not hasattr(dispatch, node.attr) ): - missing.append( - f"{path.relative_to(target_dir)}:{node.lineno}: dispatch.{node.attr}" - ) - assert not missing, missing + stale.append(f"{path.relative_to(target_dir)}:{node.lineno}: dispatch.{node.attr}") + assert not stale, stale # ------------------------------------------------------------------ the fused index score + top-16 select @@ -499,14 +558,15 @@ class _FssConfig: local_blocks = 1 -def _fss_layer(): - """What ``M3Attention._fused_score_select`` reads of its layer.""" - return SimpleNamespace(attn=SimpleNamespace(m3_config=_FssConfig()), layer_idx=5) +def _fss_core(): + """What ``DecodeTarget._fused_score_select`` reads of the core: layer 5's MSA backend configuration.""" + return SimpleNamespace(attn={5: SimpleNamespace(m3_config=_FssConfig())}) -def _fss_metadata(rows: int, qlen: int, max_tiles: int = 8192, n_valid: bool = True): - """A decode span's metadata as the fused score + select reads it (CPU stand-ins, never launched).""" - return SimpleNamespace( +def _fss_step(rows: int, qlen: int, max_tiles: int = 8192, n_valid: bool = True): + """A decode span's ``step_args`` entries as the fused score + select reads them (CPU stand-ins, never + launched).""" + return dict( msa_idx_k_cache=lambda layer_idx: torch.zeros(4, 1, 128, 128, dtype=torch.float8_e4m3fn), msa_worst_case_max_k_tiles=max_tiles, msa_n_valid_blocks=torch.ones(rows * qlen, dtype=torch.int32) if n_valid else None, @@ -517,8 +577,8 @@ def _fss_metadata(rows: int, qlen: int, max_tiles: int = 8192, n_valid: bool = T def _fss(rows: int, qlen: int = 4, dtype=torch.float8_e4m3fn, **md): idx_q = torch.zeros(rows * qlen, 128, dtype=dtype) - return layers.M3Attention._fused_score_select( - _fss_layer(), idx_q, _fss_metadata(rows, qlen, **md), rows, qlen, rows * qlen + return dispatch.DecodeTarget(_fss_core())._fused_score_select( + 5, idx_q, _fss_step(rows, qlen, **md), rows, qlen, rows * qlen ) @@ -593,13 +653,22 @@ def _lowm_operands(nk, m: int): return weight, torch.zeros(m, k, dtype=torch.bfloat16) -def _draft_linear(weight, x, bias=None): - return m3_draft_lowm.draft_linear(weight, x, bias) +_IN_TREE = "the in-tree BF16 GEMM" + + +@pytest.fixture(autouse=True) +def _in_tree_gemm(monkeypatch): + """The in-tree BF16 GEMM a drafter linear falls back to, stubbed (nothing here launches a kernel).""" + monkeypatch.setattr(dispatch, "m3_bf16_linear", lambda x, weight: _IN_TREE) + + +def _draft_linear(weight, x): + return dispatch._draft_linear(x, weight) @pytest.mark.parametrize("nk", DRAFTER_LINEARS, ids=str) def test_draft_linear_without_its_prebuilt_kernel_raises_under_require(monkeypatch, nk) -> None: - monkeypatch.setattr(m3_draft_lowm, "_BUILT", set()) + monkeypatch.setattr(dispatch, "_DRAFT_LOWM_BUILT", set()) weight, x = _lowm_operands(nk, 1) dispatch.set_require(True) with pytest.raises( @@ -611,10 +680,10 @@ def test_draft_linear_without_its_prebuilt_kernel_raises_under_require(monkeypat def test_draft_linear_without_its_prebuilt_kernel_falls_back_and_logs_once( monkeypatch, log ) -> None: - monkeypatch.setattr(m3_draft_lowm, "_BUILT", set()) + monkeypatch.setattr(dispatch, "_DRAFT_LOWM_BUILT", set()) weight, x = _lowm_operands((6144, 1024), 2) for _ in range(2): - assert _draft_linear(weight, x) is None + assert _draft_linear(weight, x) is _IN_TREE assert len(log.warning) == 1 and "m3_draft_lowm at 2 rows" in log.warning[0] @@ -632,45 +701,41 @@ def test_draft_linear_on_a_strided_activation_keeps_the_in_tree_gemm(log) -> Non wide[:, :1024] = x assert not wide[:, :1024].is_contiguous() dispatch.set_require(True) - assert _draft_linear(weight, wide[:, :1024]) is None + assert _draft_linear(weight, wide[:, :1024]) is _IN_TREE assert len(log.info) == 1 and not log.warning def test_draft_linear_past_its_rows_is_uncovered_in_every_mode(monkeypatch, log) -> None: - monkeypatch.setattr(m3_draft_lowm, "_BUILT", set()) + monkeypatch.setattr(dispatch, "_DRAFT_LOWM_BUILT", set()) dispatch.set_require(True) weight, x = _lowm_operands((6144, 1024), 3) # the o_proj table ends at 2 rows - assert _draft_linear(weight, x) is None + assert _draft_linear(weight, x) is _IN_TREE weight, x = _lowm_operands((1280, 12288), 9) - assert _draft_linear(weight, x) is None + assert _draft_linear(weight, x) is _IN_TREE weight, x = _lowm_operands((4096, 1024), 1) # not a drafter shape - assert _draft_linear(weight, x) is None + assert _draft_linear(weight, x) is _IN_TREE assert len(log.info) == 3 and not log.warning -def test_draft_linear_with_a_bias_is_not_taken(monkeypatch, log) -> None: - monkeypatch.setattr(m3_draft_lowm, "_BUILT", set()) +def test_draft_linear_on_a_batched_activation_is_uncovered(monkeypatch, log) -> None: + monkeypatch.setattr(dispatch, "_DRAFT_LOWM_BUILT", set()) dispatch.set_require(True) weight, x = _lowm_operands((6144, 1024), 1) - assert _draft_linear(weight, x, torch.zeros(6144, dtype=torch.bfloat16)) is None - assert _draft_linear(weight, x[None]) is None - assert not log.info and not log.warning + assert _draft_linear(weight, x[None]) is _IN_TREE + assert len(log.info) == 1 and not log.warning def test_prebuild_drafter_linears_caps_the_drafter_attention(monkeypatch) -> None: - """The drafter attention gets ``DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH`` and every drafter linear goes to the - low-M method before warmup; without a drafter nothing is built.""" - installed = [] - monkeypatch.setattr(m3_draft_lowm, "install", lambda dm: installed.append(dm) or 5) - monkeypatch.setattr(m3_draft_lowm, "prebuild", lambda dm: 17) - worker = spec_worker.M3Eagle3Worker.__new__(spec_worker.M3Eagle3Worker) - assert worker.prebuild_drafter_linears(None) == 0 and not installed + """The drafter attention gets ``DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH`` and every covered drafter linear call is + compiled before warmup; without a drafter nothing is built.""" + built = [] + monkeypatch.setattr(dispatch, "_prebuild_draft_lowm", lambda dm: built.append(dm) or 17) + worker = dispatch.M3Eagle3Worker.__new__(dispatch.M3Eagle3Worker) + assert worker.prebuild_drafter_linears(None) == 0 and not built attn = SimpleNamespace() - drafter = SimpleNamespace( - model=SimpleNamespace(midlayer=SimpleNamespace(self_attn=SimpleNamespace(attn=attn))) - ) + drafter = SimpleNamespace(attn=attn) assert worker.prebuild_drafter_linears(drafter) == 17 - assert installed == [drafter] + assert built == [drafter] assert attn.dense_decode_max_seq_len_by_batch is dispatch.DRAFT_DECODE_MAX_SEQ_LEN_BY_BATCH @@ -681,23 +746,8 @@ def test_prebuild_drafter_linears_caps_the_drafter_attention(monkeypatch) -> Non def _drafter(head_weight=None): - """A drafter with only what M3Eagle3Worker._m3_drafter_supported reads: one plain Eagle3 layer with a gated MLP, - a whole-table embedding, all-reduced o_proj and down_proj and a vocab-sharded BF16 head (TP8).""" - mlp = GatedMLP.__new__(GatedMLP) - torch.nn.Module.__init__(mlp) - mlp.down_proj = SimpleNamespace(all_reduce=object()) - layer = SimpleNamespace( - _next_layer_regular=False, - mlp=mlp, - self_attn=SimpleNamespace(o_proj=SimpleNamespace(all_reduce=object())), - ) - model = SimpleNamespace( - num_layers=1, - _use_mla=False, - _eh_proj_before_attn=False, - midlayer=layer, - embed_tokens=SimpleNamespace(tp_size=1), - ) + """A drafter with only what M3Eagle3Worker._m3_drafter_supported reads: its vocab-sharded BF16 head at TP8 (the + drafter's one-layer shape is validated when ``M3Eagle3Draft`` is built).""" weight = ( head_weight if head_weight is not None @@ -706,13 +756,13 @@ def _drafter(head_weight=None): lm_head = SimpleNamespace( weight=weight, bias=None, tp_mode=TensorParallelMode.COLUMN, padding_size=0, tp_size=8 ) - return SimpleNamespace(model=model, lm_head=lm_head) + return SimpleNamespace(lm_head=lm_head) def _worker(): """An M3Eagle3Worker past the step's verify: the fused verify left the first-pass gather rows and the fused first draft inputs; nothing else of the in-tree worker is read.""" - worker = spec_worker.M3Eagle3Worker.__new__(spec_worker.M3Eagle3Worker) + worker = dispatch.M3Eagle3Worker.__new__(dispatch.M3Eagle3Worker) worker.mapping = SimpleNamespace(tp_size=8, tp_rank=1, enable_attention_dp=False) worker._d2t = None worker._m3_draft_head = None @@ -734,9 +784,7 @@ def no_gpu_draft_head(monkeypatch): def test_draft_loop_without_the_prebuilt_head_raises_under_require(no_gpu_draft_head) -> None: dispatch.set_require(True) - with pytest.raises( - dispatch.FastPathUnavailable, match=r"'m3_draft_step'.* 4 rows.*prebuild_draft_head" - ): + with pytest.raises(dispatch.FastPathUnavailable, match=r"'m3_draft_step'.* 4 rows.*prebuild"): _worker()._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, _drafter()) @@ -821,9 +869,7 @@ def test_an_unsupported_drafter_is_uncovered(no_gpu_draft_head, log) -> None: worker = _worker() padded = _drafter() padded.lm_head.padding_size = 16 - two_layers = _drafter() - two_layers.model.num_layers = 2 - for drafter in (padded, two_layers): + for drafter in (padded,): assert not worker.prebuild_draft_head(drafter) assert not worker._m3_fused_loop_ok(_ATTN_MD, _SPEC_MD, 4, drafter) assert log.info and not log.warning @@ -856,7 +902,7 @@ def maybe_capture_hidden_states(self, layer_id, hidden_states, residual) -> None def _fc_early_worker(monkeypatch, eligible: bool = True): - worker = spec_worker.M3Eagle3Worker.__new__(spec_worker.M3Eagle3Worker) + worker = dispatch.M3Eagle3Worker.__new__(dispatch.M3Eagle3Worker) worker._m3_capture_norm = None worker._m3_fc_norm_w = None worker._m3_fc_early_drafter = None @@ -883,8 +929,8 @@ def test_no_early_fork_at_tp8(monkeypatch, nreq, eligible) -> None: assert worker._m3_fc_early_drafter is None for layer_id in md.layers_to_capture: worker.capture_hidden_states(md, layer_id, None) - worker.capture_hidden_states(md, spec_worker.FINAL_HIDDEN_LAYER_ID, None) - assert md.captured == [1, 29, 56, spec_worker.FINAL_HIDDEN_LAYER_ID] + worker.capture_hidden_states(md, dispatch.FINAL_HIDDEN_LAYER_ID, None) + assert md.captured == [1, 29, 56, dispatch.FINAL_HIDDEN_LAYER_ID] assert not forks and worker._m3_fc_early_stream is None @@ -893,9 +939,9 @@ def test_no_early_fork_with_a_final_state_capture(monkeypatch) -> None: stock place instead of forking at the last captured decoder layer.""" worker, forks = _fc_early_worker(monkeypatch) md = _CaptureMetadata(1) - md.layers_to_capture = (spec_worker.FINAL_HIDDEN_LAYER_ID, 29, 56) + md.layers_to_capture = (dispatch.FINAL_HIDDEN_LAYER_ID, 29, 56) worker.arm_capture_norm(SimpleNamespace(num_seqs=1), md, object(), object()) - for layer_id in (29, 56, spec_worker.FINAL_HIDDEN_LAYER_ID): + for layer_id in (29, 56, dispatch.FINAL_HIDDEN_LAYER_ID): worker.capture_hidden_states(md, layer_id, None) assert not forks and worker._m3_fc_early_stream is None @@ -903,93 +949,115 @@ def test_no_early_fork_with_a_final_state_capture(monkeypatch) -> None: # ------------------------------------------------------------------ the NVFP4 producer's coverage +def _producer_core(nvfp4: bool = True, heads: int = 8): + """What the producer dispatch reads off the core before its first precondition.""" + return SimpleNamespace(main_kv_is_nvfp4=nvfp4, heads=heads, kv_heads=1, index_heads=1) + + def test_the_producer_covers_the_tp8_nvfp4_layers_only() -> None: - """``dispatch.PRODUCER`` covers a sparse layer on an NVFP4 cache at the TP8 heads (8 query, one KV and one - index head); elsewhere the built-in producers run without a miss.""" - covered = layers.producer_pw_covered - assert covered(True, 8, 1, 1) - assert not covered(False, 8, 1, 1) - assert not covered(True, 16, 1, 1) - assert not covered(True, 8, 2, 2) + """``PRODUCER`` covers a sparse layer on an NVFP4 cache at the TP8 heads (8 query, one KV and one index head); + elsewhere the generic producer runs without a miss.""" + covered = dispatch._producer_covered + assert covered(_producer_core()) + assert not covered(_producer_core(nvfp4=False)) + assert not covered(_producer_core(heads=16)) + assert not covered(SimpleNamespace(main_kv_is_nvfp4=True, heads=8, kv_heads=2, index_heads=2)) assert dispatch.PRODUCER.keys == frozenset({"nvfp4_tp8"}) assert dispatch.PRODUCER.covers("nvfp4_tp8") assert not dispatch.PRODUCER.covers("nvfp4_tp4") -class _FakeMsaAttn: - indexer_kv_dtype = "fp8" - - -def _producer_layer(monkeypatch, nvfp4: bool = True, heads: int = 8): - """What the producer dispatch reads before its first precondition: a fused-projection MSA sparse layer.""" - monkeypatch.setattr(layers, "MiniMaxM3MsaSparseAttention", _FakeMsaAttn) - return SimpleNamespace( - enable_fused_qkv_index_projection=True, - attn=_FakeMsaAttn(), - main_kv_is_fp8=not nvfp4, - main_kv_is_nvfp4=nvfp4, - num_heads=heads, - num_key_value_heads=1, - sparse_num_index_heads=1, - pos_embd_params=None, - ) - - -def _produce(layer): +def _produce(core): # The packed TP8 projection: 8 query heads, K, V, index Q and index K of 128 each. packed = torch.zeros(4, 12 * 128, dtype=torch.bfloat16) - return layers.M3Attention._fused_fp8_qkv_indexer_norm_rope_kv_insert( - layer, packed, None, SimpleNamespace() - ) + return dispatch._sparse_qkv_producer(core, 3, packed, None, None, None, None) -def test_a_covered_producer_layer_missing_a_precondition_raises_under_require(monkeypatch) -> None: +def test_a_covered_producer_layer_missing_a_precondition_raises_under_require() -> None: """A TP8 NVFP4 sparse layer whose step lacks a producer precondition (here the position ids) is a miss.""" dispatch.set_require(True) with pytest.raises( dispatch.FastPathUnavailable, match=r"'m3_producer' covers this step at 4 rows" ): - _produce(_producer_layer(monkeypatch)) + _produce(_producer_core()) -def test_a_covered_producer_layer_without_the_fused_projection_raises_under_require( - monkeypatch, -) -> None: - layer = _producer_layer(monkeypatch) - layer.enable_fused_qkv_index_projection = False - dispatch.set_require(True) - with pytest.raises(dispatch.FastPathUnavailable, match="fuse_qkv_index_projection is disabled"): - layers.M3Attention._check_producer_projection(layer) +def _runtime(fused: bool, indexer: str = "fp8", kv: str = "fp8"): + """A model config with the runtime fields ``_validate_runtime`` reads.""" + sparse = SimpleNamespace( + algorithm="minimax_m3", + implementation="msa", + fuse_qkv_index_projection=fused, + indexer_kv_dtype=indexer, + ) + mode = SimpleNamespace( + has_fp8_kv_cache=lambda: kv == "fp8", has_fp4_kv_cache=lambda: kv == "nvfp4" + ) + return SimpleNamespace( + sparse_attention_config=sparse, quant_config=SimpleNamespace(quant_mode=mode) + ) + + +@pytest.mark.parametrize("fused", [True, False], ids=["five_way", "unfused"]) +def test_the_target_takes_both_projection_layouts(fused) -> None: + """The packed five-way projection and the separate QKV + BF16 index projections + (fuse_qkv_index_projection=False) are both this target's; a BF16 indexer or a BF16 KV cache is not.""" + dispatch._validate_runtime(_runtime(fused)) + dispatch._validate_runtime(_runtime(fused, kv="nvfp4")) + with pytest.raises(ValueError, match="indexer_kv_dtype='fp8'"): + dispatch._validate_runtime(_runtime(fused, indexer="bf16")) + with pytest.raises(ValueError, match="FP8 or NVFP4 KV cache"): + dispatch._validate_runtime(_runtime(fused, kv="bf16")) -def test_an_uncovered_layer_without_the_fused_projection_is_not_a_miss(monkeypatch) -> None: - layer = _producer_layer(monkeypatch, nvfp4=False) - layer.enable_fused_qkv_index_projection = False +@pytest.mark.parametrize("target", ["PrefillTarget", "DecodeTarget"]) +def test_an_unfused_layer_the_producer_covers_is_a_miss(target) -> None: + """As the built-in target: a layer the NVFP4 producer covers needs the five-way projection, so the unfused path is + a miss there (raised under require) before anything runs, in either target.""" dispatch.set_require(True) - layers.M3Attention._check_producer_projection(layer) + attention = getattr(dispatch, target)(_producer_core())._unfused_sparse_attention + with pytest.raises(dispatch.FastPathUnavailable, match="fuse_qkv_index_projection is disabled"): + if target == "PrefillTarget": + attention(3, None, None, None) + else: + attention(3, None, None, None, None) -def test_a_covered_producer_miss_falls_back_and_logs_once(monkeypatch, log) -> None: - layer = _producer_layer(monkeypatch) - assert _produce(layer) is None and _produce(layer) is None +def test_a_covered_producer_miss_falls_back_and_logs_once(log) -> None: + core = _producer_core() + assert _produce(core) is None and _produce(core) is None assert len(log.warning) == 1 and "m3_producer at 4 rows" in log.warning[0] @pytest.mark.parametrize("nvfp4,heads", [(False, 8), (True, 16)], ids=["fp8_cache", "tp4_heads"]) -def test_an_uncovered_producer_layer_falls_back_silently(monkeypatch, log, nvfp4, heads) -> None: +def test_an_uncovered_producer_layer_falls_back_silently(log, nvfp4, heads) -> None: dispatch.set_require(True) - assert _produce(_producer_layer(monkeypatch, nvfp4, heads)) is None + assert _produce(_producer_core(nvfp4, heads)) is None assert not log.warning def test_a_covered_producer_layer_with_the_wrong_pool_layout_is_not_supported() -> None: - """The pool-layout check the covered branch takes before the native fallback (then a miss).""" + """The pool-layout check the covered branch takes before the generic producer (then a miss).""" pool = torch.zeros(4, 2, 1, 128, 32, dtype=torch.uint8) scales = torch.zeros(4, 2, 1, 128, 8, dtype=torch.uint8) - assert not layers.producer_pw_supported(8, 1, 1, pool, scales) + assert not dispatch._producer_layout(pool, scales) good = torch.zeros(4, 2, 1, 128, 64, dtype=torch.uint8) - assert layers.producer_pw_supported(8, 1, 1, good, scales) - assert not layers.producer_pw_supported(16, 1, 1, good, scales) - assert not layers.producer_pw_supported( - 8, 1, 1, good, scales.transpose(3, 4).contiguous().transpose(3, 4) - ) + assert dispatch._producer_layout(good, scales) + assert not dispatch._producer_layout(good, scales.transpose(3, 4).contiguous().transpose(3, 4)) + + +# ------------------------------------------------------------------ the shell's speculative plumbing + + +def test_the_shell_installs_guided_decoding_on_the_worker() -> None: + """With speculation the guided decoder goes to the worker (the in-tree one-engine shell's delegation); without a + worker the shell reports it cannot take one.""" + shell_cls = dispatch.ModelingV2MiniMaxM3Nvfp4Sm103Tp8 + shell = shell_cls.__new__(shell_cls) + torch.nn.Module.__init__(shell) + installed = [] + shell.spec_worker = SimpleNamespace(set_guided_decoder=lambda g: installed.append(g) or True) + decoder = object() + assert shell.set_guided_decoder(decoder) and installed == [decoder] + shell.spec_worker = None + assert shell.set_guided_decoder(decoder) is False diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_prebuild.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_prebuild.py index 580f419a3f7f..e932865571e5 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_prebuild.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_prebuild.py @@ -1,28 +1,26 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The decode kernels ``build_targets`` builds before the executor's warmup (``prebuild_kernels``), and that a forward -after it never compiles one: the MoE engine of every ``MOE_ENGINE`` token count, the decode GEMVs with their push -variants, and the index scorer's prefetching kernel at every decode query length (the tp8 target never selects the -scorer's early-issue schedule: ``INDEX_SCORE_EARLY_MAX_REQUESTS`` is 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 engine of every ``MOE_ENGINE`` token count, the decode GEMVs with their push variants, and +the index scorer's prefetching kernel at every decode query length (the tp8 target never selects the scorer's +early-issue schedule: ``INDEX_SCORE_EARLY_MAX_REQUESTS`` is 0). 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 -from unittest import mock import pytest import torch -from torch import nn 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_103__tp8 import ( # noqa: E501 - dispatch, - layers, - linears, modeling, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 + m3_index_score as index_score_kernel, +) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels.mx_gemv import ( # noqa: E501 op as mx_gemv, ) @@ -35,11 +33,23 @@ QKV_SPARSE, ) +dispatch = modeling + assert torch.cuda.is_available(), "the prebuild checks require a CUDA device" INDEX_HEADS = 1 # index heads per rank at TP8 INDEX_DIM = 128 DRAFT_LEN = 3 +_CORE_METHODS = ( + "prebuild_kernels", + "_prebuild_engines", + "_try_build", + "engine", + "_lookup", + "_scorer_variants", + "_prebuild_index_scorer", + "decode_query_lens", +) @pytest.fixture(autouse=True) @@ -57,56 +67,38 @@ def __init__(self, indexer_kv_dtype: str): self.m3_config = SimpleNamespace(num_index_heads=INDEX_HEADS, sparse_index_dim=INDEX_DIM) -def _moe(layer_idx: int, shared: bool = True) -> layers.M3MoE: - """An M3MoE built by its own constructor over a stub stock block: only what prebuild_engines and the lookups - read (the builds are stubbed).""" - - def stock_init(self, model_config, aux_stream_dict, layer_idx=None): - nn.Module.__init__(self) - # One rank: the build agreement takes no collective. - self.mapping = SimpleNamespace(tp_size=1, enable_attention_dp=False) - self.shared_experts = None - - with mock.patch.object(layers.MiniMaxM3MoE, "__init__", stock_init): - moe = layers.M3MoE(None, {}, layer_idx) - moe.shared_experts = object() if shared else None - return moe - - -def _model(monkeypatch, num_layers: int = 3, indexer_kv_dtype: str = "fp8"): - """The parts of MiniMaxM3Tp8ModelingV2 prebuild_kernels reads: layers (MoE + sparse attention), the embedding - device and the speculative config.""" - monkeypatch.setattr(layers, "MiniMaxM3MsaSparseAttention", _FakeSparseAttn) - model_layers = [] - for idx in range(num_layers): - model_layers.append( - SimpleNamespace( - block_sparse_moe=_moe(idx, shared=idx != 1) if idx else None, - self_attn=SimpleNamespace( - is_sparse_attention_layer=idx != 0, attn=_FakeSparseAttn(indexer_kv_dtype) - ), - ) - ) - stub = SimpleNamespace( - layers=model_layers, - embed_tokens=SimpleNamespace(weight=torch.empty(1, device="cuda")), +def _core(num_layers: int = 3, indexer_kv_dtype: str = "fp8", built=None, fail=False): + """The parts of ``MiniMaxM3Tp8ModelingV2`` 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), + 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={}, + projections=lambda: [], ) - stub.decode_query_lens = lambda: modeling.MiniMaxM3Tp8ModelingV2.decode_query_lens(stub) - return stub + for name in _CORE_METHODS: + setattr(core, name, getattr(modeling.MiniMaxM3Tp8ModelingV2, 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 _stub_engine_builds(monkeypatch, built: list) -> None: - def build(self, num_tokens, shared=True): - built.append(("moe", self.layer_idx, num_tokens, shared)) - return ("engine", self.layer_idx, num_tokens) + core._build_engine = build + return core - monkeypatch.setattr(layers.M3MoE, "_build_engine", build) - -def _forbid_compiles(monkeypatch) -> None: +def _forbid_compiles(monkeypatch, core=None) -> None: """Every compile / build entry point a forward could reach raises.""" import cutlass.cute as cute @@ -116,112 +108,110 @@ def boom(*args, **kwargs): monkeypatch.setattr(cute, "compile", boom) monkeypatch.setattr(mx_gemv, "ensure_compiled", boom) monkeypatch.setattr(mx_gemv, "precompile", boom) - monkeypatch.setattr(layers.M3MoE, "_build_engine", boom) - monkeypatch.setattr(layers._scorer_runner(), "_compile", boom) + monkeypatch.setattr(index_score_kernel.IndexScoreRunner, "_compile", boom) + if core is not None: + core._build_engine = boom -def test_decode_query_lens_follow_the_speculative_config(monkeypatch) -> None: - model = _model(monkeypatch) - assert model.decode_query_lens() == [1, 1 + DRAFT_LEN] - model.model_config.spec_config = None - assert model.decode_query_lens() == [1] - model.model_config.spec_config = SimpleNamespace( +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 model.decode_query_lens() == [1, 3, 5] + assert core.decode_query_lens() == [1, 3, 5] -def test_prebuild_builds_every_engine_count_and_scorer_variant(monkeypatch) -> None: +def test_prebuild_builds_every_engine_count_and_scorer_variant() -> None: built = [] - _stub_engine_builds(monkeypatch, built) - monkeypatch.setattr(modeling, "adopt_m3_linears", lambda root: built.append(("linears",))) - model = _model(monkeypatch) - times = modeling.MiniMaxM3Tp8ModelingV2.prebuild_kernels(model) + core = _core(built=built) + times = core.prebuild_kernels() - assert built[0] == ("linears",) assert mx_gemv.frozen(), "prebuild must close the GEMV compiles" - # Layer 0 is dense, layer 1 has no shared expert (no engine takes it); layer 2 gets every count, in order. - assert built[1:] == [("moe", 2, n, True) for n in range(1, dispatch.MOE_ENGINE_MAX_TOKENS + 1)] - moe = model.layers[2].block_sparse_moe - assert set(moe._engines) == {(n, True, True) for n in dispatch.MOE_ENGINE.keys} - assert not hasattr(moe, "_moe16_engines") - runner = layers._scorer_runner() + # Layer 0 is dense; every MoE layer gets every count, in layer order. + def per_layer(i): + return [("moe", i, n, True) for n in range(1, dispatch.MOE_ENGINE_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 not hasattr(core, "moe16_engines") + want = { - (torch.float8_e4m3fn, INDEX_HEADS, qlen, INDEX_DIM, layers.INDEX_SCORE_SPLIT_K, False) + (torch.float8_e4m3fn, INDEX_HEADS, qlen, INDEX_DIM, modeling.INDEX_SCORE_SPLIT_K, False) for qlen in (1, 1 + DRAFT_LEN) } assert set(times["scorer_keys"]) == want - assert want <= set(runner.kernel_cache) + assert want <= set(index_score_kernel.IndexScoreRunner.kernel_cache) for key in ("gemv_s", "moe_s", "scorer_s"): assert times[key] >= 0.0 - # A second build_targets builds nothing again. + # A second build builds no engine again. built.clear() - modeling.MiniMaxM3Tp8ModelingV2.prebuild_kernels(model) - assert built == [("linears",)] + core.prebuild_kernels() + assert built == [] def test_engine_lookups_after_prebuild_never_build(monkeypatch) -> None: - built = [] - _stub_engine_builds(monkeypatch, built) - moe = _moe(5) - moe.prebuild_engines() - _forbid_compiles(monkeypatch) + 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 moe.engine(n, shared=True) == ("engine", 5, n) + assert core.engine(5, n, shared=True) == ("engine", 5, n) def test_missing_engine_is_a_miss_not_a_build(monkeypatch) -> None: - _forbid_compiles(monkeypatch) - moe = _moe(7) + core = _core() + core.moe_engines[7] = {} + _forbid_compiles(monkeypatch, core) dispatch.set_require(False) - assert moe.engine(3, shared=True) is None + assert core.engine(7, 3, shared=True) is None dispatch.set_require(True) with pytest.raises(dispatch.FastPathUnavailable, match="not built before warmup"): - moe.engine(3, shared=True) + core.engine(7, 3, shared=True) -def test_failed_engine_build_is_logged_then_missed(monkeypatch) -> None: - def fail(self, num_tokens, shared=True): - raise ValueError("experts backend is not TRTLLM-Gen NVFP4 SwiGLU-OAI") - - monkeypatch.setattr(layers.M3MoE, "_build_engine", fail) - moe = _moe(9) - moe.prebuild_engines() - assert all(v is None for v in moe._engines.values()) and len(moe._engines) == 8 +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"): - moe.engine(1, shared=True) + 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, m3_push=push, tp_size=8) + return SimpleNamespace(weight=wq.view(n, k), weight_scale=wsf, push=push) @pytest.mark.parametrize( "nk,push", - ((O_PROJ, linears.PUSH_ATTN), (QKV_SPARSE, None), (DENSE_DOWN, None)), + ((O_PROJ, "attn"), (QKV_SPARSE, None), (DENSE_DOWN, None)), ids=("o_proj", "qkv_sparse", "dense_down"), ) def test_gemvs_after_prebuild_never_compile(monkeypatch, nk, push) -> None: use_pdl = True module = _projection(nk, push) - variants = linears.gemv_variants(module) + variants = modeling._gemv_variants(*nk, push, 8) 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 == linears.PUSH_ATTN else dispatch.MOE_PUSH.keys + 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) - assert linears.prebuild_m3_gemvs(module, use_pdl) > 0 + # 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) @@ -269,13 +259,11 @@ def test_frozen_gemv_declines_an_uncompiled_kernel(monkeypatch) -> None: @pytest.mark.parametrize("indexer_kv_dtype", ("fp8", "bf16")) def test_scorer_after_prebuild_never_compiles(monkeypatch, indexer_kv_dtype) -> None: - model = _model(monkeypatch, indexer_kv_dtype=indexer_kv_dtype) - keys = layers.prebuild_index_scorer( - model.layers, model.decode_query_lens(), torch.device("cuda") - ) + core = _core(indexer_kv_dtype=indexer_kv_dtype) + keys = core._prebuild_index_scorer(core.decode_query_lens(), torch.device("cuda")) # The decode query lengths 1 and 1 + DRAFT_LEN, the prefetching kernel only. assert len(keys) == 2 and not any(key[-1] for key in keys) - _forbid_compiles(monkeypatch) + _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 @@ -298,8 +286,3 @@ def test_scorer_after_prebuild_never_compiles(monkeypatch, indexer_kv_dtype) -> early, ) torch.cuda.synchronize() - - -def test_moe_keeps_its_layer_index(): - # prebuild_engines names the layer in its build and agreement messages. - assert _moe(5).layer_idx == 5 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py index f0b7fcd51ecf..ee3c95ac63be 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py @@ -10,11 +10,9 @@ import pytest import torch -from torch import nn from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 - dispatch, - layers, + modeling, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 m3_mc_allreduce as mc, @@ -23,10 +21,10 @@ @pytest.fixture(autouse=True) def _restore_modes(): - require = dispatch.require_mode() + require = modeling.require_mode() yield - dispatch.set_require(require) - dispatch.reset_logged() + modeling.set_require(require) + modeling.reset_logged() @pytest.mark.parametrize("finalize", [False, True]) @@ -133,25 +131,21 @@ def test_tp_agree_built_two_ranks(two_rank_comm, oks): assert out == (built if all(oks) else [None, None]) -def _moe(layer_idx: int, fail: bool, world: int = 2) -> layers.M3MoE: - """An M3MoE on ``world`` simulated TP ranks whose engine builds are stubbed (and fail where ``fail``).""" +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 stock_init(self, model_config, aux_stream_dict, layer_idx=None): - nn.Module.__init__(self) - self.mapping = _mapping(world) - self.shared_experts = None - - with mock.patch.object(layers.MiniMaxM3MoE, "__init__", stock_init): - moe = layers.M3MoE(None, {}, layer_idx) - moe.shared_experts = object() - - def build(num_tokens, shared=True): + def build(i, num_tokens, shared=True): if fail: raise RuntimeError("compile failed on this rank") return ("engine", num_tokens) - moe._build_engine = build - return moe + 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( @@ -161,29 +155,26 @@ def build(num_tokens, shared=True): ) def test_engine_build_agreed_across_ranks(monkeypatch, world, failing): comm = _patch_comm(monkeypatch, world) - moes = [_moe(3, rank in failing, world) for rank in range(world)] - _run_ranks([m.prebuild_engines for m in moes]) - keys = sorted(moes[0]._engines) + 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(m, "_moe16_engines") for m in moes) - for m in moes[1:]: - assert m._engines.keys() == moes[0]._engines.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 m in moes: - assert (m._engines[key] is None) == bool(failing), key - # Every rank then takes the same path in the forward: the engine, or (after a miss) the built-in MoE. - dispatch.set_require(False) - for n in sorted(dispatch.MOE_ENGINE.keys): - got = [m.engine(n, shared=True) is None for m in moes] + 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): - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.spec_worker import ( # noqa: E501 - M3Eagle3Worker, - ) - + 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) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py index 8fa2ee6f331f..8bc830b733ef 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The tp8 target's files that are byte copies of the tp4 target's (and of the shared CuTe DSL GEMV kernel) stay so. +"""The tp8 target's files that are byte copies of the tp4 target's stay so. The two MiniMax-M3 targets share no code by design (``test_modeling_v2_claims.test_targets_do_not_share_files``): the tp8 target was made as a copy of the tp4 one and keeps its own copy of every file, tp8-specific or not. The files @@ -15,7 +15,6 @@ import pytest -import tensorrt_llm._torch as _torch_pkg import tensorrt_llm._torch._experimental.modeling_v2 as _modeling_v2 # Nothing of either target is imported (no GPU, no built extensions needed): the paths are the packages', resolved @@ -27,8 +26,14 @@ _TP8 = _M3 / "minimax_m3_nvfp4__sm_103__tp8" #: Paths relative to the target directory that are byte copies of the tp4 target's. +#: (``ops/attention/m3_msa_attention.py`` left the list: its reference dequantizes NVFP4 pages with each target's own +#: long-context sparse decode entry, ``m3_sparse_decode_lc`` at TP4 and ``m3_sparse_decode_lc8`` at TP8.) IDENTICAL_TO_TP4 = ( + "kernels/embedding_firstnorm/__init__.py", + "kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py", + "kernels/embedding_firstnorm/op.py", "kernels/long_select.py", + "kernels/m3_force_vec.py", "kernels/m3_index_score.py", "kernels/m3_prequant.py", "kernels/m3_router_gemv.py", @@ -36,15 +41,23 @@ "kernels/m3_shard_accept.py", "kernels/m3_spec_tail.py", "kernels/msa_kv_update.py", + "kernels/mx_gemv/m3_mx_gemv_kernel.py", "kernels/triton_sparse_decode_mq.py", "ops/__init__.py", + "ops/activation/__init__.py", "ops/attention/__init__.py", + "ops/attention/m3_draft_attention.py", + "ops/attention/m3_fp8_indexer_qk_norm_rope.py", + "ops/attention/m3_select_blocks.py", "ops/attention/msa_kv_update.py", "ops/comm/__init__.py", + "ops/embedding/__init__.py", + "ops/embedding/m3_embedding.py", "ops/gemm/__init__.py", + "ops/gemm/m3_fp32_linear.py", "ops/moe/__init__.py", - "ops/moe/m3_router_gemv.py", "ops/moe/m3_router_gemv_fp4.py", + "ops/moe/m3_router_gemv.py", "ops/moe/m3_router_tc.py", "ops/norm/__init__.py", "ops/norm/m3_embedding_firstnorm.py", @@ -52,29 +65,14 @@ "ops/spec/m3_capture_norm.py", "ops/spec/m3_chunk_norm.py", "ops/spec/m3_draft_input.py", + "ops/spec/m3_force_vec.py", "ops/spec/m3_slot_scatter.py", "ops/spec/m3_step_gather.py", ) -#: The tp8 target's own copy of the CuTe DSL decode GEMV kernel (its op.py carries the TP8 tables), against the shared -#: one the tp4 target uses. -GEMV_KERNEL = "kernels/mx_gemv/m3_mx_gemv_kernel.py" -_SHARED_GEMV_KERNEL = ( - Path(_torch_pkg.__file__).resolve().parent - / "cute_dsl_kernels" - / "minimax_m3" - / "mx_gemv" - / "m3_mx_gemv_kernel.py" -) - def test_the_targets_are_where_this_test_reads_them() -> None: - for path in ( - _TP4 / "dispatch.py", - _TP8 / "dispatch.py", - _TP8 / GEMV_KERNEL, - _SHARED_GEMV_KERNEL, - ): + for path in (_TP4 / "modeling.py", _TP8 / "modeling.py"): assert path.is_file(), path @@ -100,14 +98,6 @@ def test_the_copy_is_the_tp4_file(relative) -> None: ) -def test_the_gemv_kernel_is_the_shared_kernel() -> None: - diff = _differs(_TP8 / GEMV_KERNEL, _SHARED_GEMV_KERNEL) - assert not diff, ( - f"{diff}. The tp8 target's decode GEMV kernel is a copy of the shared CuTe DSL one: a change to one copy " - "must be considered for the other." - ) - - def test_the_list_names_every_identical_file() -> None: """A file that became identical again (or was copied unchanged later) is listed, so it is guarded too.""" unlisted = sorted( 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 8b15df66b4aa..7acd64d04e67 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -310,6 +310,284 @@ def test_minimax_m3_parallel_names_the_target_topology(mapping_kwargs, parallel) assert routing._parallel(Mapping(**mapping_kwargs)) == parallel +# The MiniMax-M3 Eagle3 drafter of the released checkpoint (inferact MiniMax-M3-EAGLE3-GQA), as its config.json reads. +_M3_DRAFTER = dict( + num_hidden_layers=1, + fc_norm=True, + hidden_size=6144, + num_attention_heads=64, + num_key_value_heads=4, + head_dim=128, + intermediate_size=18432, + vocab_size=200064, + draft_vocab_size=200064, + attention_bias=False, + tie_word_embeddings=False, +) + + +def _m3_drafter_spec(tmp_path, captured=None, **overrides): + import json + from types import SimpleNamespace + + (tmp_path / "config.json").write_text(json.dumps({**_M3_DRAFTER, **overrides})) + return SimpleNamespace(speculative_model=str(tmp_path), eagle3_layers_to_capture=captured) + + +@pytest.mark.parametrize( + "overrides, captured, reason", + [ + ({}, None, None), + ({}, {1, 29, 56}, None), + (dict(num_hidden_layers=2), None, "2 drafter layers, not 1"), + (dict(fc_norm=False), None, "no per-layer fc_norm"), + (dict(eagle_config={"eh_proj_before_attn": True}), None, "eh_proj before attention"), + (dict(eagle_config={"norm_before_fc": True}), None, "norm before the fc"), + (dict(kv_lora_rank=512), None, "an MLA drafter layer"), + (dict(target_hidden_size=4096), None, "a target hidden size other than the drafter's"), + (dict(draft_vocab_size=32000), None, "a draft vocabulary other than the target's"), + (dict(mlp_bias=True), None, "biased projections"), + (dict(tie_word_embeddings=True), None, "a tied head"), + ({}, {-1}, "a single captured target layer"), + (dict(num_key_value_heads=1), None, "num_key_value_heads=1, not 4"), + (dict(intermediate_size=14336), None, "intermediate_size=14336, not 18432"), + ( + dict(hidden_size=4096, head_dim=None), + None, + "hidden_size=4096, not 6144; head_dim=64, not 128", + ), + ], + ids=[ + "released_drafter", + "three_captured_layers", + "two_layers", + "no_fc_norm", + "eh_proj_before_attn", + "norm_before_fc", + "mla", + "target_hidden", + "draft_vocab", + "bias", + "tied_head", + "one_captured_layer", + "one_kv_head", + "mlp_width", + "hidden_size", + ], +) +def test_minimax_m3_routes_only_the_drafter_its_targets_build( + tmp_path, overrides, captured, reason +): + """The targets draft with their own one-layer fc_norm Eagle3 drafter: any other drafter is declined in routing + (``auto`` takes the built-in model, ``require`` names the criterion) instead of failing at construction.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + spec = _m3_drafter_spec(tmp_path, captured, **overrides) + assert routing._drafter_mismatch(spec) == reason + assert routing._drafter_mismatch(None) is None + + +@pytest.mark.parametrize("drafter_layers", [1, 3], ids=["flat_drafter", "multi_layer_drafter"]) +def test_minimax_m3_route_declines_a_drafter_it_does_not_build( + tmp_path, monkeypatch, drafter_layers +): + 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 + from tensorrt_llm.quantization.mode import QuantAlgo + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + monkeypatch.setattr(routing, "_eagle3_one_model_or_none", lambda spec: True) + text_config = dict( + zip( + ( + "num_hidden_layers", + "hidden_size", + "num_local_experts", + "intermediate_size", + "num_attention_heads", + "num_key_value_heads", + ), + routing._SHAPE, + ) + ) + spec = _m3_drafter_spec(tmp_path, num_hidden_layers=drafter_layers) + spec.spec_dec_mode = SimpleNamespace(name="EAGLE3_ONE_MODEL") + mapping = Mapping(world_size=4, tp_size=4, moe_tp_size=4, moe_ep_size=1) + pretrained = SimpleNamespace(text_config=text_config) + runtime = _m3_runtime() + if drafter_layers == 1: + target = routing.route(pretrained, mapping, routing._SM, spec, runtime=runtime) + assert target.__name__ == "ModelingV2MiniMaxM3Nvfp4Sm103Tp4" + else: + with pytest.raises(NoTarget, match="Eagle3 drafter: 3 drafter layers, not 1"): + routing.route(pretrained, mapping, routing._SM, spec, runtime=runtime) + + +def _m3_sparse(**overrides): + """The sparse-attention runtime of the PR deployment (MSA, FP8 indexer, fused projection), with overrides.""" + from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig + + fields = dict( + implementation="msa", + indexer_kv_dtype="fp8", + fuse_qkv_index_projection=True, + sparse_disable_index_value=True, + ) + fields.update(overrides) + return MiniMaxM3SparseAttentionConfig(**fields) + + +def _m3_runtime(sparse=None, kv="NVFP4"): + from tensorrt_llm._torch._experimental.modeling_v2._router_index import RuntimeFacts + + return RuntimeFacts( + sparse_attention_config=_m3_sparse() if sparse is None else sparse, kv_cache_quant_algo=kv + ) + + +# Every runtime the engine accepts for MiniMax-M3 that the targets are not written for; the built-in model runs them. +_M3_RUNTIME_DECLINES = [ + ( + dict(implementation="triton", indexer_kv_dtype="bf16", fuse_qkv_index_projection=False), + "FP8", + "sparse attention implementation 'triton', not 'msa'; indexer_kv_dtype 'bf16', not 'fp8'", + ), + ( + dict(indexer_kv_dtype="bf16", fuse_qkv_index_projection=False), + "FP8", + "indexer_kv_dtype 'bf16', not 'fp8'", + ), + ( + dict(fuse_qkv_index_projection=False), + None, + "a BF16 main KV cache, not FP8 or NVFP4", + ), +] +_M3_RUNTIME_DECLINE_IDS = ["triton", "bf16_indexer", "bf16_main_kv"] + + +@pytest.mark.parametrize( + "sparse, kv", + [ + (dict(), "NVFP4"), + (dict(), "FP8"), + (dict(fuse_qkv_index_projection=False), "FP8"), + (dict(fuse_qkv_index_projection=False), "NVFP4"), + ], + ids=["fused_nvfp4", "fused_fp8", "unfused_fp8", "unfused_nvfp4"], +) +def test_minimax_m3_runtime_the_targets_are_written_for(sparse, kv): + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + assert routing._runtime_mismatch(_m3_runtime(_m3_sparse(**sparse), kv)) is None + + +@pytest.mark.parametrize("sparse, kv, reason", _M3_RUNTIME_DECLINES, ids=_M3_RUNTIME_DECLINE_IDS) +def test_minimax_m3_runtime_mismatch_names_the_setting(sparse, kv, reason): + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + assert routing._runtime_mismatch(_m3_runtime(_m3_sparse(**sparse), kv)) == reason + + +def test_minimax_m3_runtime_without_its_sparse_config_is_declined(): + from tensorrt_llm._torch._experimental.modeling_v2._router_index import RuntimeFacts + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + runtime = RuntimeFacts(sparse_attention_config=None, kv_cache_quant_algo="NVFP4") + assert routing._runtime_mismatch(runtime) == "no MiniMax-M3 sparse_attention_config" + + +def _m3_resolve_config(sparse, kv_checkpoint, tp=4): + """A ModelConfig as the loader holds it when it resolves: the MiniMax-M3 VL wrapper with the target text shape, + the checkpoint's KV quantization (before kv_cache_config.dtype is applied) and the sparse-attention runtime.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + from tensorrt_llm.models.modeling_utils import QuantConfig + + text_config = dict( + zip( + ( + "num_hidden_layers", + "hidden_size", + "num_local_experts", + "intermediate_size", + "num_attention_heads", + "num_key_value_heads", + ), + routing._SHAPE, + ) + ) + pretrained = PretrainedConfig( + architectures=["MiniMaxM3SparseForConditionalGeneration"], text_config=text_config + ) + return ModelConfig( + pretrained_config=pretrained, + mapping=Mapping(world_size=tp, tp_size=tp, moe_tp_size=tp, moe_ep_size=1), + quant_config=QuantConfig(kv_cache_quant_algo=kv_checkpoint), + sparse_attention_config=sparse, + ) + + +def _m3_llm_args(mode, kv_dtype): + return _llm_args(mode, kv_cache_config=dict(dtype=kv_dtype)) + + +@pytest.mark.parametrize("tp", [4, 8]) +@pytest.mark.parametrize("mode", ["auto", "require"]) +def test_minimax_m3_pr_runtime_resolves_to_the_target(monkeypatch, mode, tp): + """The PR deployment (MSA, FP8 indexer, fused projection, kv_cache_config.dtype='nvfp4') takes the target.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + from tensorrt_llm.quantization.mode import QuantAlgo + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + config = _m3_resolve_config(_m3_sparse(), None, tp) + target = modeling_v2_resolve(config, _m3_llm_args(mode, "nvfp4")) + assert target.__name__ == f"ModelingV2MiniMaxM3Nvfp4Sm103Tp{tp}" + + +@pytest.mark.parametrize("sparse, kv, reason", _M3_RUNTIME_DECLINES, ids=_M3_RUNTIME_DECLINE_IDS) +@pytest.mark.parametrize("tp", [4, 8]) +def test_minimax_m3_other_runtimes_fall_back_under_auto(monkeypatch, sparse, kv, reason, tp): + """``auto`` resolves nothing, so the loader builds the registered built-in MiniMax-M3 model.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + from tensorrt_llm._torch.models.modeling_minimaxm3 import MiniMaxM3VLForConditionalGeneration + from tensorrt_llm.quantization.mode import QuantAlgo + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + # kv_cache_config.dtype 'auto' keeps the checkpoint's KV quantization (None: the BF16 cache). + config = _m3_resolve_config(_m3_sparse(**sparse), kv, tp) + assert modeling_v2_resolve(config, _m3_llm_args("auto", "auto")) is None + resolved = replace(config, modeling_v2_target=None) + assert AutoModelForCausalLM._resolve_class(resolved) is MiniMaxM3VLForConditionalGeneration + + +@pytest.mark.parametrize("sparse, kv, reason", _M3_RUNTIME_DECLINES, ids=_M3_RUNTIME_DECLINE_IDS) +def test_minimax_m3_other_runtimes_raise_under_require(monkeypatch, sparse, kv, reason): + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + from tensorrt_llm.quantization.mode import QuantAlgo + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + config = _m3_resolve_config(_m3_sparse(**sparse), kv) + with pytest.raises(ValueError) as excinfo: + modeling_v2_resolve(config, _m3_llm_args("require", "auto")) + message = str(excinfo.value) + assert "MiniMaxM3SparseForConditionalGeneration" in message + assert f"runtime: {reason}" in message + + +def test_minimax_m3_kv_cache_dtype_overrides_the_checkpoint(monkeypatch): + """An explicit kv_cache_config.dtype is the KV cache the engine runs with, whatever the checkpoint declares.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + from tensorrt_llm.quantization.mode import QuantAlgo + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + config = _m3_resolve_config(_m3_sparse(fuse_qkv_index_projection=False), None) + assert modeling_v2_resolve(config, _m3_llm_args("auto", "auto")) is None + target = modeling_v2_resolve(config, _m3_llm_args("auto", "fp8")) + assert target.__name__ == "ModelingV2MiniMaxM3Nvfp4Sm103Tp4" + + # --- qwen4_exp: one-model MTP-Eagle over a linear chain, one rank --------------- _QWEN4_EXP_TARGET = "models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.modeling." From dcdbd8c0cf2f54f368963476b9acb0db98cfddb3 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:40:28 -0700 Subject: [PATCH 07/14] [None][chore] MiniMax-M3: keep the M3 torch ops, kernels and fake impls in M3 folders Review follow-up, no logic change: - cpp/tensorrt_llm/kernels/stepInputsScatter.{cu,h} move to kernels/minimaxM3/. - The MiniMax-M3 torch ops (thop/minimaxM3*Op.cpp and stepInputsScatterOp.cpp) move to thop/minimaxM3/, which adds them to th_common_attention with target_sources() (as thop/moe/ does for th_common_moe). - The fake impls of the ops the decode targets add (step_inputs_scatter, minimax_m3_mxfp8_low_m_gemm, minimax_m3_router_gemv / _gemv_fp4 / _tc, minimax_m3_producer_pw, minimax_m3_mc_allreduce_norm) move from custom_ops/cpp_custom_ops.py to modeling_v2/models/minimax_m3/fake_ops.py, imported by the MiniMax-M3 modeling_v2 package, by pyexecutor/step_inputs.py and by the fake-registration test. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../{ => minimaxM3}/stepInputsScatter.cu | 2 +- .../{ => minimaxM3}/stepInputsScatter.h | 0 cpp/tensorrt_llm/thop/CMakeLists.txt | 24 +---- .../thop/minimaxM3/CMakeLists.txt | 19 ++++ .../{ => minimaxM3}/minimaxM3AttnC1Op.cpp | 0 .../{ => minimaxM3}/minimaxM3AttnMegaOp.cpp | 0 .../{ => minimaxM3}/minimaxM3DraftStepOp.cpp | 0 .../{ => minimaxM3}/minimaxM3Fp8IndexerOp.cpp | 0 .../minimaxM3FusedDecodeAttnOp.cpp | 0 .../minimaxM3McAllreduceNormOp.cpp | 0 .../minimaxM3MsaKvUpdateOp.cpp | 0 .../minimaxM3Mxfp8LowMGemmOp.cpp | 0 .../{ => minimaxM3}/minimaxM3ProducerPwOp.cpp | 0 .../{ => minimaxM3}/minimaxM3RouterGemvOp.cpp | 0 .../minimaxM3SampledAcceptOp.cpp | 0 .../minimaxM3SelectBlocksOp.cpp | 0 .../{ => minimaxM3}/minimaxM3SelectLongOp.cpp | 0 .../minimaxM3ShardAcceptOp.cpp | 0 .../minimaxM3SparseDecodeLcOp.cpp | 0 .../{ => minimaxM3}/minimaxM3SpecTailOp.cpp | 0 .../{ => minimaxM3}/minimaxM3TargetHeadOp.cpp | 0 .../minimaxM3Tp8FusedScoreSelectOp.cpp | 0 .../minimaxM3Tp8SparseDecodeLc8Op.cpp | 0 .../{ => minimaxM3}/stepInputsScatterOp.cpp | 8 +- .../modeling_v2/models/minimax_m3/__init__.py | 2 + .../modeling_v2/models/minimax_m3/fake_ops.py | 101 ++++++++++++++++++ .../_torch/custom_ops/cpp_custom_ops.py | 54 ---------- tensorrt_llm/_torch/pyexecutor/step_inputs.py | 3 + .../_torch/executor/test_step_inputs.py | 2 +- .../parallel_hw_agnostic/test_custom_ops.py | 3 + 30 files changed, 138 insertions(+), 80 deletions(-) rename cpp/tensorrt_llm/kernels/{ => minimaxM3}/stepInputsScatter.cu (97%) rename cpp/tensorrt_llm/kernels/{ => minimaxM3}/stepInputsScatter.h (100%) create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/CMakeLists.txt rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3AttnC1Op.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3AttnMegaOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3DraftStepOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3Fp8IndexerOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3FusedDecodeAttnOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3McAllreduceNormOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3MsaKvUpdateOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3Mxfp8LowMGemmOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3ProducerPwOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3RouterGemvOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3SampledAcceptOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3SelectBlocksOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3SelectLongOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3ShardAcceptOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3SparseDecodeLcOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3SpecTailOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3TargetHeadOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3Tp8FusedScoreSelectOp.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/minimaxM3Tp8SparseDecodeLc8Op.cpp (100%) rename cpp/tensorrt_llm/thop/{ => minimaxM3}/stepInputsScatterOp.cpp (87%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/fake_ops.py diff --git a/cpp/tensorrt_llm/kernels/stepInputsScatter.cu b/cpp/tensorrt_llm/kernels/minimaxM3/stepInputsScatter.cu similarity index 97% rename from cpp/tensorrt_llm/kernels/stepInputsScatter.cu rename to cpp/tensorrt_llm/kernels/minimaxM3/stepInputsScatter.cu index 31616af5add0..2e80de17bfd9 100644 --- a/cpp/tensorrt_llm/kernels/stepInputsScatter.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3/stepInputsScatter.cu @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/stepInputsScatter.h" +#include "tensorrt_llm/kernels/minimaxM3/stepInputsScatter.h" TRTLLM_NAMESPACE_BEGIN diff --git a/cpp/tensorrt_llm/kernels/stepInputsScatter.h b/cpp/tensorrt_llm/kernels/minimaxM3/stepInputsScatter.h similarity index 100% rename from cpp/tensorrt_llm/kernels/stepInputsScatter.h rename to cpp/tensorrt_llm/kernels/minimaxM3/stepInputsScatter.h diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 16d646338f32..dcf4871276da 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -67,26 +67,6 @@ add_th_common_family( IndexerKCacheScatterOp.cpp IndexerTopKOp.cpp sparseKvCacheCompactOp.cpp - minimaxM3AttnC1Op.cpp - minimaxM3AttnMegaOp.cpp - minimaxM3DraftStepOp.cpp - minimaxM3Fp8IndexerOp.cpp - minimaxM3FusedDecodeAttnOp.cpp - minimaxM3McAllreduceNormOp.cpp - minimaxM3MsaKvUpdateOp.cpp - minimaxM3Mxfp8LowMGemmOp.cpp - minimaxM3ProducerPwOp.cpp - minimaxM3RouterGemvOp.cpp - minimaxM3SampledAcceptOp.cpp - minimaxM3SelectBlocksOp.cpp - minimaxM3SelectLongOp.cpp - minimaxM3ShardAcceptOp.cpp - minimaxM3SparseDecodeLcOp.cpp - minimaxM3SpecTailOp.cpp - minimaxM3TargetHeadOp.cpp - minimaxM3Tp8FusedScoreSelectOp.cpp - minimaxM3Tp8SparseDecodeLc8Op.cpp - stepInputsScatterOp.cpp mlaRopeInplaceOp.cpp relativeAttentionBiasOp.cpp dsv3RopeOp.cpp @@ -94,6 +74,10 @@ add_th_common_family( nvfp4MlaKvCacheGatherOp.cpp trtllmGenQKVProcessOp.cpp) +# MiniMax-M3 torch ops live in minimaxM3/; that directory adds them to +# th_common_attention with target_sources() so the list stays next to the files. +add_subdirectory(minimaxM3) + add_th_common_family( th_common_gemm_quant cutlassScaledMM.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3/CMakeLists.txt b/cpp/tensorrt_llm/thop/minimaxM3/CMakeLists.txt new file mode 100644 index 000000000000..0ea4ebdce2e0 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/CMakeLists.txt @@ -0,0 +1,19 @@ +# 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. + +# torch.ops.trtllm MiniMax-M3 ops (kernels in kernels/minimaxM3/). The parent +# defines th_common_attention. +file(GLOB SRC_CPP CONFIGURE_DEPENDS *.cpp) +target_sources(th_common_attention PRIVATE ${SRC_CPP}) diff --git a/cpp/tensorrt_llm/thop/minimaxM3AttnC1Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3AttnC1Op.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3AttnC1Op.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3AttnC1Op.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3AttnMegaOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3AttnMegaOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3AttnMegaOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3AttnMegaOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3DraftStepOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3DraftStepOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3DraftStepOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3DraftStepOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3Fp8IndexerOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Fp8IndexerOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3Fp8IndexerOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Fp8IndexerOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3FusedDecodeAttnOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3FusedDecodeAttnOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3FusedDecodeAttnOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3FusedDecodeAttnOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3McAllreduceNormOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3McAllreduceNormOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3McAllreduceNormOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3McAllreduceNormOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3MsaKvUpdateOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3MsaKvUpdateOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3MsaKvUpdateOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3MsaKvUpdateOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3Mxfp8LowMGemmOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Mxfp8LowMGemmOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3Mxfp8LowMGemmOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Mxfp8LowMGemmOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3ProducerPwOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3ProducerPwOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3ProducerPwOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3ProducerPwOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3RouterGemvOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3RouterGemvOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3RouterGemvOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3RouterGemvOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3SampledAcceptOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SampledAcceptOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3SampledAcceptOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SampledAcceptOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SelectBlocksOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SelectBlocksOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3SelectLongOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SelectLongOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3SelectLongOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SelectLongOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3ShardAcceptOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3ShardAcceptOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3ShardAcceptOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3ShardAcceptOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3SparseDecodeLcOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SparseDecodeLcOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3SparseDecodeLcOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SparseDecodeLcOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3SpecTailOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SpecTailOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3SpecTailOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3SpecTailOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3TargetHeadOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3TargetHeadOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3TargetHeadOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3TargetHeadOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3Tp8FusedScoreSelectOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Tp8FusedScoreSelectOp.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3Tp8FusedScoreSelectOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Tp8FusedScoreSelectOp.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3Tp8SparseDecodeLc8Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Tp8SparseDecodeLc8Op.cpp similarity index 100% rename from cpp/tensorrt_llm/thop/minimaxM3Tp8SparseDecodeLc8Op.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Tp8SparseDecodeLc8Op.cpp diff --git a/cpp/tensorrt_llm/thop/stepInputsScatterOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/stepInputsScatterOp.cpp similarity index 87% rename from cpp/tensorrt_llm/thop/stepInputsScatterOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/stepInputsScatterOp.cpp index e8d168d68849..b78a779a29d1 100644 --- a/cpp/tensorrt_llm/thop/stepInputsScatterOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/stepInputsScatterOp.cpp @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/stepInputsScatter.h" +#include "tensorrt_llm/kernels/minimaxM3/stepInputsScatter.h" #include #include @@ -28,9 +28,9 @@ namespace torch_ext //! Issue a batch of packed host-to-device and device-to-device copies in one launch. //! -//! @param stage Contiguous uint8 CUDA tensor laid out as described in kernels/stepInputsScatter.h. Its destinations are -//! raw device addresses: the caller guarantees each one is live and that no other work on the stream writes it between -//! staging and this call. +//! @param stage Contiguous uint8 CUDA tensor laid out as described in kernels/minimaxM3/stepInputsScatter.h. Its +//! destinations are raw device addresses: the caller guarantees each one is live and that no other work on the stream +//! writes it between staging and this call. void stepInputsScatter(torch::Tensor const& stage) { TORCH_CHECK(stage.is_cuda(), "step_inputs_scatter expects a CUDA staging buffer"); diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/__init__.py index a3dfe4a6ce39..5c445343d34d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/__init__.py @@ -1,3 +1,5 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """MiniMaxM3SparseForConditionalGeneration targets.""" + +from . import fake_ops # noqa: F401 (registers the fake impls of the targets' torch ops) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/fake_ops.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/fake_ops.py new file mode 100644 index 000000000000..c47398888520 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/fake_ops.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fake (meta) implementations of the torch ops that the MiniMax-M3 decode targets add (``trtllm::step_inputs_scatter`` +and the ``trtllm::minimax_m3_*`` decode ops in ``cpp/tensorrt_llm/thop/minimaxM3/``), for torch.compile and fake-tensor +tracing. Importing this module registers them once; the MiniMax-M3 modeling_v2 package and the step-input pack import +it.""" + +import torch + + +@torch.library.register_fake("trtllm::step_inputs_scatter") +def _(stage): + # Writes the staged rows into their destinations; returns nothing. + return None + + +@torch.library.register_fake("trtllm::minimax_m3_mxfp8_low_m_gemm") +def _(x, w, wsf, warps, cluster_k, row_tiles, pdl): + return x.new_empty((x.shape[0], w.shape[0])) + + +@torch.library.register_fake("trtllm::minimax_m3_router_gemv") +def _(x, w, rows, pdl): + return x.new_empty((x.shape[0], w.shape[0]), dtype=torch.float32) + + +def _minimax_m3_router_fp4_outputs(x, w): + m, k = x.shape + return [ + x.new_empty((m, w.shape[0]), dtype=torch.float32), + x.new_empty((m, k // 2), dtype=torch.uint8), + x.new_empty((m * (k // 16),), dtype=torch.uint8), + ] + + +@torch.library.register_fake("trtllm::minimax_m3_router_gemv_fp4") +def _(x, w, sf_scale, pdl): + return _minimax_m3_router_fp4_outputs(x, w) + + +@torch.library.register_fake("trtllm::minimax_m3_router_tc") +def _(x, w, sf_scale, pdl): + outputs = _minimax_m3_router_fp4_outputs(x, w) + return outputs if sf_scale is not None else outputs[:1] + + +@torch.library.register_fake("trtllm::minimax_m3_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, + heads_q, + pdl, + early, +): + num_tokens = packed.shape[0] + return [ + packed.new_empty((num_tokens, heads_q, 128), dtype=torch.float8_e4m3fn), + packed.new_empty((num_tokens, 1, 128), dtype=torch.float8_e4m3fn), + ] + + +@torch.library.register_fake("trtllm::minimax_m3_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, +): + 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 diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index b7993fe66b8f..18b68a373e21 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -418,60 +418,6 @@ def _( return scores.new_empty((scores.shape[2], scores.shape[0], topk), dtype=torch.int32) - @torch.library.register_fake("trtllm::step_inputs_scatter") - def _(stage): - # Writes the staged rows into their destinations; returns nothing. - return None - - @torch.library.register_fake("trtllm::minimax_m3_mxfp8_low_m_gemm") - def _(x, w, wsf, warps, cluster_k, row_tiles, pdl): - return x.new_empty((x.shape[0], w.shape[0])) - - @torch.library.register_fake("trtllm::minimax_m3_router_gemv") - def _(x, w, rows, pdl): - return x.new_empty((x.shape[0], w.shape[0]), dtype=torch.float32) - - def _minimax_m3_router_fp4_outputs(x, w): - m, k = x.shape - return [ - x.new_empty((m, w.shape[0]), dtype=torch.float32), - x.new_empty((m, k // 2), dtype=torch.uint8), - x.new_empty((m * (k // 16), ), dtype=torch.uint8), - ] - - @torch.library.register_fake("trtllm::minimax_m3_router_gemv_fp4") - def _(x, w, sf_scale, pdl): - return _minimax_m3_router_fp4_outputs(x, w) - - @torch.library.register_fake("trtllm::minimax_m3_router_tc") - def _(x, w, sf_scale, pdl): - outputs = _minimax_m3_router_fp4_outputs(x, w) - return outputs if sf_scale is not None else outputs[:1] - - @torch.library.register_fake("trtllm::minimax_m3_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, heads_q, pdl, early): - num_tokens = packed.shape[0] - return [ - packed.new_empty((num_tokens, heads_q, 128), - dtype=torch.float8_e4m3fn), - packed.new_empty((num_tokens, 1, 128), dtype=torch.float8_e4m3fn), - ] - - @torch.library.register_fake("trtllm::minimax_m3_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): - 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 - @torch.library.register_fake("trtllm::kda_decode") def _(x_q: torch.Tensor, x_k: torch.Tensor, x_v: torch.Tensor, w_q_t: torch.Tensor, w_k_t: torch.Tensor, w_v_t: torch.Tensor, diff --git a/tensorrt_llm/_torch/pyexecutor/step_inputs.py b/tensorrt_llm/_torch/pyexecutor/step_inputs.py index 2f47dadaba8d..ff0c02a95611 100644 --- a/tensorrt_llm/_torch/pyexecutor/step_inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/step_inputs.py @@ -29,6 +29,9 @@ import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import ( # noqa: F401 + fake_ops as _m3_fake_ops, # registers trtllm::step_inputs_scatter's fake impl +) from tensorrt_llm._utils import prefer_pinned diff --git a/tests/unittest/_torch/executor/test_step_inputs.py b/tests/unittest/_torch/executor/test_step_inputs.py index f7f3859a5d05..c1b4b95a02d4 100644 --- a/tests/unittest/_torch/executor/test_step_inputs.py +++ b/tests/unittest/_torch/executor/test_step_inputs.py @@ -21,7 +21,7 @@ def _scatter_op_registered() -> bool: pytest.mark.skipif( not _scatter_op_registered(), reason="torch.ops.trtllm.step_inputs_scatter is not registered (built from " - "cpp/tensorrt_llm/thop/stepInputsScatterOp.cpp)", + "cpp/tensorrt_llm/thop/minimaxM3/stepInputsScatterOp.cpp)", ), ] diff --git a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_custom_ops.py b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_custom_ops.py index 6b0878159b37..2d21af5db11f 100644 --- a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_custom_ops.py +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_custom_ops.py @@ -19,6 +19,9 @@ import tensorrt_llm # noqa: F401 +# The MiniMax-M3 ops' fake impls are registered by their modeling_v2 package. +import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.fake_ops # noqa: F401 + pytestmark = pytest.mark.cpu_only From dd6a5b710a5c485d013829dfb153d110aece0ddc Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:23:57 -0700 Subject: [PATCH 08/14] [None][fix] Decoder runner warmup: read the model's runtime hooks through runtime_hooks_of The autotuner and CUDA-graph warmup passes read merge_tuned_tactics_across_ranks from the runner's cached hooks, which only the constructor sets; engines built another way (main's warmup tests drive DecoderRunner._run_autotuner_warmup on a stub) failed with an AttributeError. They now read runtime_hooks_of(self.model), the all-off default for a model that declares no hooks, the same value the constructor caches. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../_torch/pyexecutor/engine/runners/decoder/runner.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 4afb50927eae..9bc88f0e8416 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -1798,7 +1798,9 @@ def run_autotuner_pass(autotune_context: Any, synchronize_trtllm_cache: bool) -> ran_flashinfer_forward = run_autotuner_pass( flashinfer_mxfp8_autotune(), synchronize_trtllm_cache=False ) - if self.runtime_hooks.merge_tuned_tactics_across_ranks: + # The model's declared hooks (the all-off default when it declares none): the warmup passes also run on + # engines built without the runner's constructor. + if runtime_hooks_of(self.model).merge_tuned_tactics_across_ranks: # Rank-local profiling noise must not leave tensor-parallel peers # serving different tactics for one shape: keep the fastest everywhere. # Collective; before the native MXFP8 cache below registers the tactics. @@ -2199,7 +2201,10 @@ def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): ) with flashinfer_autotune_context, flashinfer_mxfp8_decode_graph_capture(): self._capture_generation_cuda_graphs(resource_manager) - if is_warmup_only and self.runtime_hooks.merge_tuned_tactics_across_ranks: + if ( + is_warmup_only + and runtime_hooks_of(self.model).merge_tuned_tactics_across_ranks + ): # Tactics tuned while warming the decode graphs agree across # ranks before the capture pass reads them. AutoTuner.get().post_tune_merge_tactics() From f1d15d645c1eb8d221a2f984c02a35bf6ebee2da Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:25:44 -0700 Subject: [PATCH 09/14] [None][fix] MiniMax-M3 KV cache manager: reused pool views follow the HND storage layout Since #20159 the dense KV storage is head-major and get_buffers hands a token-major request a transposed view. The reused pool views (reuse_pool_views) cached the untransposed view and transposed it on every call, so a reused view was a new tensor each time. The transpose is now part of the view's cache key and applied before caching. The runtime-policy tests expect the view in the manager's own layout over its physical one. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../sparse/minimax_m3/cache_manager.py | 18 +++++++++++++++--- .../sparse/msa/test_msa_runtime_policy.py | 19 +++++++++++++------ 2 files changed, 28 insertions(+), 9 deletions(-) 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 ef275f46bdf6..6182057757a3 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 @@ -848,10 +848,17 @@ def get_buffers( addr_key, torch_dtype, num_slots, scale, page_shape = self._kv_slot_geometry( layer_idx, self._main_kv_layout_name() ) - kv = self._pool_view("kv", addr_key, torch_dtype, num_slots, scale, page_shape) # Sparse reference kernels use token-major indexing; preserve the # native head-major backing and expose a zero-copy strided view. - return kv if kv_layout == self._main_kv_layout_name() else kv.transpose(2, 3) + return self._pool_view( + "kv", + addr_key, + torch_dtype, + num_slots, + scale, + page_shape, + transpose=kv_layout != self._main_kv_layout_name(), + ) def configure_decode_shortcuts( self, skip_unchanged_block_offsets: bool, reuse_pool_views: bool @@ -874,8 +881,10 @@ def _pool_view( num_slots: int, scale: int, page_shape: List[int], + transpose: bool = False, ) -> torch.Tensor: - """The ``[num_slots, 2, ...]`` K+V view of a pool rooted at ``addr_key``. + """The ``[num_slots, 2, ...]`` K+V view of a pool rooted at ``addr_key``; ``transpose`` swaps the + page's head and token axes (the token-major view of head-major storage). The view is a pure function of the pool geometry, which is static once the pool is laid out, so with ``reuse_pool_views`` it is built once per geometry @@ -889,6 +898,7 @@ def _pool_view( num_slots, scale, tuple(page_shape), + transpose, torch.is_inference_mode_enabled(), ) view = self._pool_views.get(key) if self.reuse_pool_views else None @@ -898,6 +908,8 @@ def _pool_view( TensorWrapper(addr_key, torch_dtype, full_slot_shape) ) view = full_view[:, :2] + if transpose: + view = view.transpose(2, 3) if self.reuse_pool_views: self._pool_views[key] = view return view 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 e9bbc98acd07..2f63cd962e39 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 @@ -304,9 +304,16 @@ def kv_scale_slot_geometry(self, layer_idx, kv_layout): return kv_pool, scale_pool -def _assert_kv_view(view: torch.Tensor, pool: torch.Tensor) -> None: +def _assert_kv_view( + view: torch.Tensor, pool: torch.Tensor, manager: MiniMaxM3KVCacheManagerV2 +) -> None: + # The pool is laid out in the manager's physical layout; a view in the other layout swaps the + # page's head and token axes over the same memory. + page = list(PAGE_SHAPE) + if manager._main_kv_view_layout != manager._main_kv_layout_name(): + page[0], page[1] = page[1], page[0] assert view.data_ptr() == pool.data_ptr() - assert list(view.shape) == [NUM_SLOTS, 2, *PAGE_SHAPE] + assert list(view.shape) == [NUM_SLOTS, 2, *page] assert view.stride(0) == SCALE * 2 * 8 * 16 @@ -319,8 +326,8 @@ def test_pool_views_are_rebuilt_on_every_call_by_default(manager, pools): assert not manager.reuse_pool_views assert first is not second - _assert_kv_view(first, kv_pool) - _assert_kv_view(second, kv_pool) + _assert_kv_view(first, kv_pool, manager) + _assert_kv_view(second, kv_pool, manager) @requires_cuda @@ -334,7 +341,7 @@ def test_reused_pool_views_are_built_once_per_geometry(manager, pools): assert manager.get_buffers(3) is view assert manager.get_block_scale_buffers(3) is scale_view assert scale_view is not view - _assert_kv_view(view, kv_pool) + _assert_kv_view(view, kv_pool, manager) assert scale_view.data_ptr() == scale_pool.data_ptr() assert scale_view.dtype == torch.uint8 @@ -352,7 +359,7 @@ def test_reused_pool_views_are_kept_per_inference_mode(manager, pools): assert inference_view is not view assert inference_view.is_inference() assert not view.is_inference() - _assert_kv_view(inference_view, kv_pool) + _assert_kv_view(inference_view, kv_pool, manager) @requires_cuda From 8ccd79f11d80d0512f8c204b784499988d116c72 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Mon, 19 Oct 2026 20:18:18 -0700 Subject: [PATCH 10/14] [None][test] Decoder runner init test: the runner config fixture declares is_spec_decode DecoderRunner reads DecoderRunnerConfig.is_spec_decode at construction (the spec step fold binds the engine buffers). The compiled-body validation test builds the runner from a SimpleNamespace config; give it the field the real config always has. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../_torch/compilation/test_phase_selective_forward.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/compilation/test_phase_selective_forward.py b/tests/unittest/_torch/compilation/test_phase_selective_forward.py index 613975c89b9b..33e5d42ab924 100644 --- a/tests/unittest/_torch/compilation/test_phase_selective_forward.py +++ b/tests/unittest/_torch/compilation/test_phase_selective_forward.py @@ -170,7 +170,9 @@ def test_decoder_runner_validates_compiled_body_at_initialization( torch.nn.Module.__init__(model) wrapper, _, _ = _create_prefill_compiled_model() model.model = wrapper if has_wrapper else torch.nn.Identity() - config = SimpleNamespace(compile_only_piecewise_graphs=enabled, max_beam_width=1) + config = SimpleNamespace( + compile_only_piecewise_graphs=enabled, max_beam_width=1, is_spec_decode=False + ) from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import runner as runner_module for component in ("DecoderInputBuffers", "DecoderMetadata", "DecoderWarmup"): From 300c1c1e6422ebb848956365486b74436487f0d3 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 21 Oct 2026 13:47:56 -0700 Subject: [PATCH 11/14] [None][chore] MiniMax-M3 modeling_v2: standard route interface, layout guard, target kernels under ops Review follow-ups, no kernel or numerics change: - The shared router calls every routing module with the standard route(pretrained_config, mapping, compute_capability, spec_config, kv_cache_quant_algo); RuntimeFacts and ROUTES_ON_RUNTIME are removed. The M3 route checks the KV cache from kv_cache_quant_algo; the sparse-attention runtime (MSA, FP8 indexer) is an engine argument, so the targets' _validate_runtime checks it at construction. - The M3 route drops the Eagle3 drafter structure checks, which the target's drafter repeats at construction, and declines a checkpoint whose RoPE base or sparse-attention layers differ from the ones the targets build (the architecture class and shape alone do not pin them). - Each sm_103 target's kernels/ moves into its ops/ directory; imports follow. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../_experimental/modeling_v2/__init__.py | 2 - .../modeling_v2/_router_index.py | 43 +-- .../kernels/__init__.py | 3 - .../minimax_m3_nvfp4__sm_103__tp4/modeling.py | 54 ++-- .../ops/attention/m3_fused_decode_attn.py | 2 +- .../ops/attention/m3_index_score.py | 4 +- .../ops/attention/m3_sparse_decode.py | 2 +- .../ops/comm/m3_mc_allreduce_norm.py | 2 +- .../embedding_firstnorm/__init__.py | 0 .../m3_embedding_firstnorm_kernel.py | 0 .../embedding_firstnorm/op.py | 0 .../ops/gemm/m3_mx_gemv.py | 2 +- .../ops/gemm/m3_mx_gemv_bf16.py | 2 +- .../{kernels => ops}/long_select.py | 0 .../{kernels => ops}/m3_attn_lc.py | 0 .../{kernels => ops}/m3_attn_mega.py | 0 .../{kernels => ops}/m3_attn_mega_c1.py | 0 .../{kernels => ops}/m3_draft_step.py | 0 .../{kernels => ops}/m3_force_vec.py | 0 .../{kernels => ops}/m3_fused_attn.py | 0 .../{kernels => ops}/m3_index_score.py | 0 .../{kernels => ops}/m3_mc_allreduce.py | 0 .../{kernels => ops}/m3_moe/__init__.py | 0 .../m3_moe/m3_moe16_kernel.py | 0 .../{kernels => ops}/m3_moe/m3_moe_kernel.py | 0 .../{kernels => ops}/m3_moe/op.py | 0 .../{kernels => ops}/m3_moe/op16.py | 0 .../{kernels => ops}/m3_prequant.py | 0 .../{kernels => ops}/m3_router_gemv.py | 0 .../{kernels => ops}/m3_sampled_accept.py | 0 .../{kernels => ops}/m3_shard_accept.py | 0 .../{kernels => ops}/m3_spec_tail.py | 0 .../{kernels => ops}/m3_target_head.py | 0 .../ops/moe/m3_moe16_engine.py | 4 +- .../ops/moe/m3_moe_engine.py | 2 +- .../{kernels => ops}/msa_kv_update.py | 0 .../{kernels => ops}/mx_gemv/__init__.py | 0 .../mx_gemv/m3_mx_gemv_kernel.py | 0 .../{kernels => ops}/mx_gemv/op.py | 0 .../{kernels => ops}/mxfp8_low_m_gemm.py | 0 .../ops/norm/m3_embedding_firstnorm.py | 2 +- .../ops/spec/m3_force_vec.py | 2 +- .../{kernels => ops}/triton_sparse_decode.py | 0 .../triton_sparse_decode_mq.py | 0 .../kernels/__init__.py | 3 - .../minimax_m3_nvfp4__sm_103__tp8/modeling.py | 40 +-- .../ops/attention/m3_index_score.py | 4 +- .../ops/attention/m3_sparse_decode.py | 2 +- .../ops/comm/m3_mc_allreduce_norm.py | 2 +- .../embedding_firstnorm/__init__.py | 0 .../m3_embedding_firstnorm_kernel.py | 0 .../embedding_firstnorm/op.py | 0 .../ops/gemm/m3_mx_gemv.py | 2 +- .../ops/gemm/m3_mx_gemv_bf16.py | 2 +- .../{kernels => ops}/long_select.py | 0 .../{kernels => ops}/m3_attn_lc8.py | 0 .../{kernels => ops}/m3_draft_step.py | 0 .../{kernels => ops}/m3_force_vec.py | 0 .../{kernels => ops}/m3_fused_score_select.py | 0 .../{kernels => ops}/m3_index_score.py | 0 .../{kernels => ops}/m3_mc_allreduce.py | 0 .../{kernels => ops}/m3_moe/__init__.py | 0 .../{kernels => ops}/m3_moe/m3_moe_kernel.py | 0 .../{kernels => ops}/m3_moe/op.py | 0 .../{kernels => ops}/m3_prequant.py | 0 .../{kernels => ops}/m3_router_gemv.py | 0 .../{kernels => ops}/m3_sampled_accept.py | 0 .../{kernels => ops}/m3_shard_accept.py | 0 .../{kernels => ops}/m3_spec_tail.py | 0 .../ops/moe/m3_moe_engine.py | 2 +- .../{kernels => ops}/msa_kv_update.py | 0 .../{kernels => ops}/mx_gemv/__init__.py | 0 .../mx_gemv/m3_mx_gemv_kernel.py | 0 .../{kernels => ops}/mx_gemv/op.py | 0 .../ops/norm/m3_embedding_firstnorm.py | 2 +- .../ops/spec/m3_force_vec.py | 2 +- .../{kernels => ops}/triton_sparse_decode.py | 0 .../triton_sparse_decode_mq.py | 0 .../modeling_v2/models/minimax_m3/routing.py | 171 +++------- .../attention/test_modeling_v2_m3_attn_c1.py | 6 +- .../test_modeling_v2_m3_attn_mega.py | 6 +- .../test_modeling_v2_m3_fused_decode_attn.py | 6 +- .../test_modeling_v2_msa_kv_update.py | 6 +- .../comm/_m3_mc_allreduce_norm_op_matrix.py | 2 +- .../gemm/test_modeling_v2_m3_target_head.py | 2 +- .../spec/test_modeling_v2_m3_draft_input.py | 2 +- .../spec/test_modeling_v2_m3_draft_step.py | 2 +- .../test_modeling_v2_m3_exact_numerics.py | 4 +- .../test_modeling_v2_m3_dispatch.py | 10 +- .../test_modeling_v2_m3_prebuild.py | 8 +- .../test_modeling_v2_m3_rank_agreement.py | 2 +- .../test_modeling_v2_m3_fused_score_select.py | 4 +- .../test_modeling_v2_msa_kv_update.py | 6 +- .../comm/_m3_mc_allreduce_norm_op_matrix.py | 2 +- .../spec/test_modeling_v2_m3_draft_input.py | 2 +- .../spec/test_modeling_v2_m3_draft_step.py | 2 +- .../test_modeling_v2_m3_exact_numerics.py | 8 +- .../test_modeling_v2_m3_dispatch.py | 20 +- .../test_modeling_v2_m3_prebuild.py | 8 +- .../test_modeling_v2_m3_rank_agreement.py | 2 +- .../test_modeling_v2_m3_tp4_drift.py | 28 +- .../modeling_v2/test_modeling_v2_routing.py | 296 +++++++----------- 102 files changed, 297 insertions(+), 495 deletions(-) delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/__init__.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/embedding_firstnorm/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/embedding_firstnorm/m3_embedding_firstnorm_kernel.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/embedding_firstnorm/op.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/long_select.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_attn_lc.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_attn_mega.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_attn_mega_c1.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_draft_step.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_force_vec.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_fused_attn.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_index_score.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_mc_allreduce.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_moe/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_moe/m3_moe16_kernel.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_moe/m3_moe_kernel.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_moe/op.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_moe/op16.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_prequant.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_router_gemv.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_sampled_accept.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_shard_accept.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_spec_tail.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/m3_target_head.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/msa_kv_update.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/mx_gemv/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/mx_gemv/m3_mx_gemv_kernel.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/mx_gemv/op.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/mxfp8_low_m_gemm.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/triton_sparse_decode.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/{kernels => ops}/triton_sparse_decode_mq.py (100%) delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/__init__.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/embedding_firstnorm/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/embedding_firstnorm/m3_embedding_firstnorm_kernel.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/embedding_firstnorm/op.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/long_select.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/m3_attn_lc8.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/m3_draft_step.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/m3_force_vec.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/m3_fused_score_select.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/m3_index_score.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/m3_mc_allreduce.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/m3_moe/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/m3_moe/m3_moe_kernel.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/m3_moe/op.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/m3_prequant.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/m3_router_gemv.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/m3_sampled_accept.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/m3_shard_accept.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/m3_spec_tail.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/msa_kv_update.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/mx_gemv/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/mx_gemv/m3_mx_gemv_kernel.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/mx_gemv/op.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/triton_sparse_decode.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/{kernels => ops}/triton_sparse_decode_mq.py (100%) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py index a9c6c10e7552..7dc4b711fd91 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py @@ -31,7 +31,6 @@ MODELING_V2_VALIDATE_ENV, ModelingV2Mode, NoTarget, - RuntimeFacts, modeling_v2_resolve, resolved_mode, step_contract_enabled, @@ -42,7 +41,6 @@ "MODELING_V2_VALIDATE_ENV", "ModelingV2Mode", "NoTarget", - "RuntimeFacts", "modeling_v2_resolve", "resolved_mode", "step_contract_enabled", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index c4f9cc1c2356..60a138b565f1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -18,9 +18,7 @@ * The routing module's ``route`` owns the rest of the decision as one forward-reading tree over the facts that fix which code runs: the checkpoint's config, the parallel mapping, the GPU's compute capability and - the speculative-decoding config (and, for a module that sets - ``ROUTES_ON_RUNTIME``, the engine arguments in :class:`RuntimeFacts`). It - returns the target class, or raises + the speculative-decoding config. It returns the target class, or raises :class:`NoTarget` with the reason in words. Reading that one file tells you where any configuration lands. @@ -38,9 +36,8 @@ import enum import os -from dataclasses import dataclass from importlib import import_module -from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple +from typing import TYPE_CHECKING, Dict, Optional, Tuple from tensorrt_llm.logger import logger @@ -108,36 +105,6 @@ class NoTarget(Exception): """ -@dataclass(frozen=True) -class RuntimeFacts: - """Engine arguments that fix which code runs but are not in the checkpoint. - - A routing module that sets ``ROUTES_ON_RUNTIME = True`` receives them as - ``route(..., runtime=RuntimeFacts)`` and declines (``NoTarget``) a runtime - its targets are not written for, so ``auto`` falls back to the built-in - implementation instead of the target refusing at construction. - """ - - # ``ModelConfig.sparse_attention_config`` (the LLM API's sparse_attention_config), or None. - sparse_attention_config: Any - # The main KV cache quantization the engine will run with (a ``QuantAlgo`` value, None for the - # unquantized cache): ``kv_cache_config.dtype``, or the checkpoint's under ``auto``. - kv_cache_quant_algo: Optional[str] - - @classmethod - def of(cls, config: "ModelConfig", llm_args: "TorchLlmArgs") -> "RuntimeFacts": - # The resolver runs before the loader applies kv_cache_config.dtype to the quant config - # (validate_and_set_kv_cache_quant); this is the value that function will set. - from tensorrt_llm._torch.pyexecutor.model_loader import _KV_CACHE_MAP - - dtype = llm_args.kv_cache_config.dtype - checkpoint = config.quant_config.kv_cache_quant_algo if config.quant_config else None - return cls( - sparse_attention_config=config.sparse_attention_config, - kv_cache_quant_algo=checkpoint if dtype == "auto" else _KV_CACHE_MAP.get(dtype), - ) - - # The mode each target class was last resolved under (the loader builds the class right after resolving it). _RESOLVED_MODES: Dict[type, ModelingV2Mode] = {} @@ -237,18 +204,12 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti raise NoTarget( f"no routing module for this architecture; routed architectures: {known}" ) - runtime = ( - {"runtime": RuntimeFacts.of(config, llm_args)} - if getattr(routing, "ROUTES_ON_RUNTIME", False) - else {} - ) target = routing.route( pretrained_config, config.mapping, _device_compute_capability(), config.spec_config, _kv_cache_quant_algo(config, llm_args), - **runtime, ) except NoTarget as why: if mode is ModelingV2Mode.REQUIRE: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/__init__.py deleted file mode 100644 index a111efce7a4a..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Kernels of the MiniMax-M3 NVFP4 / sm_103 / tp4 target.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py index e627ac7a7419..5fe483d48cc7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py @@ -40,7 +40,7 @@ **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 fused -decode attention for the o_proj. It publishes it in ``kernels/m3_prequant.SLOT`` keyed by the BF16 tensor; the next +decode attention 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. @@ -145,23 +145,16 @@ from tensorrt_llm.quantization.mode import QuantAlgo from . import weights as _weights -from .kernels import m3_attn_mega as _attn_mega_kernel -from .kernels import m3_attn_mega_c1 as _attn_c1_kernel -from .kernels import m3_draft_step as m3_draft_step_kernel -from .kernels import m3_fused_attn as _fused_attn_kernel -from .kernels import m3_index_score as _index_score_kernel -from .kernels import m3_mc_allreduce, m3_prequant, m3_spec_tail -from .kernels import m3_sampled_accept as m3_sampled_accept_kernel -from .kernels import m3_shard_accept as m3_shard_accept_kernel -from .kernels import m3_target_head as m3_target_head_kernel -from .kernels import msa_kv_update as _msa_kv_update_kernel -from .kernels.long_select import long_select_blocks -from .kernels.m3_router_gemv import m3_router_gemv as _router_gemv -from .kernels.m3_router_gemv import m3_router_gemv_fp4 as _router_gemv_fp4 -from .kernels.m3_router_gemv import m3_router_tc as _router_tc -from .kernels.mx_gemv import op as mx_gemv -from .kernels.mxfp8_low_m_gemm import mxfp8_low_m_gemm as _low_m_gemm -from .kernels.triton_sparse_decode import MERGE_MXFP8, MERGE_MXFP8_MIN_TOKENS +from .ops import m3_attn_mega as _attn_mega_kernel +from .ops import m3_attn_mega_c1 as _attn_c1_kernel +from .ops import m3_draft_step as m3_draft_step_kernel +from .ops import m3_fused_attn as _fused_attn_kernel +from .ops import m3_index_score as _index_score_kernel +from .ops import m3_mc_allreduce, m3_prequant, m3_spec_tail +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 m3_target_head as m3_target_head_kernel +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_attn_c1 import m3_attn_c1 from .ops.attention.m3_attn_mega import m3_attn_mega @@ -201,17 +194,24 @@ 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.gemm.mxfp8_low_m_gemm import mxfp8_low_m_gemm +from .ops.long_select import long_select_blocks +from .ops.m3_router_gemv import m3_router_gemv as _router_gemv +from .ops.m3_router_gemv import m3_router_gemv_fp4 as _router_gemv_fp4 +from .ops.m3_router_gemv import m3_router_tc as _router_tc 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_gemv import m3_router_gemv from .ops.moe.m3_router_gemv_fp4 import m3_router_gemv_fp4 from .ops.moe.m3_router_tc import m3_router_tc +from .ops.mx_gemv import op as mx_gemv +from .ops.mxfp8_low_m_gemm import mxfp8_low_m_gemm as _low_m_gemm 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_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_shard_accept import m3_shard_accept +from .ops.triton_sparse_decode import MERGE_MXFP8, MERGE_MXFP8_MIN_TOKENS _SM = (10, 3) # -------------------------------------------------------------------------------------------------------------------- @@ -238,7 +238,7 @@ class DecodePlan: # 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 (kernels/m3_mc_allreduce.BUCKETS). +# verify; verify: the c4 / c8 verifies. Each is one multicast all-reduce bucket (ops/m3_mc_allreduce.BUCKETS). DECODE_PLANS = ( DecodePlan("tiny", 8, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), DecodePlan("small", 16, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), @@ -286,22 +286,22 @@ class DecodePlan: # One request verifying a draft of 6 query rows runs attention + o_proj as one launch (96 CTAs: 6 clusters of 16). ATTN_C1_QUERY_LEN = 6 # capacity: fixed: the verify kernel's query rows # 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 (kernels/m3_shard_accept), at the configured draft length. +# 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 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 (kernels/m3_target_head; the same BF16 x BF16 -> fp32 +# 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. TARGET_HEAD_MAX_ROWS = TARGET_HEAD_KERNEL_MAX_ROWS # 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 = 8 # capacity: fixed: measured fork crossover # 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 (kernels/m3_spec_tail.draft_input), the drafter layer's in-tree attention, MLP and +# 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 (kernels/m3_draft_step), at the +# 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. @@ -309,7 +309,7 @@ class DecodePlan: DRAFT_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN # capacity: fixed: the built draft length DRAFT_STEP_MAX_REQUESTS = DRAFT_STEP_KERNEL_MAX_TOKENS # The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows -# (kernels/triton_sparse_decode.MQ_MIN_TOKENS). +# (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. @@ -350,7 +350,7 @@ class DecodePlan: DENSE_DOWN: 16, } # The fused low-M MXFP8 GEMM (projections marked for it), per projection: its built configurations' row limit -# (kernels/mxfp8_low_m_gemm). The bf16 GEMV serves every one of these rows first; the GEMM takes them only if it +# (ops/mxfp8_low_m_gemm). The bf16 GEMV serves every one of these rows first; the GEMM takes them only if it # declines. LOW_M_GEMM_MAX_TOKENS: Dict[Tuple[int, int], int] = { QKV_SPARSE: 8, @@ -815,8 +815,8 @@ def _text_model_config(model_config: ModelConfig) -> ModelConfig: def _validate_runtime(model_config: ModelConfig) -> None: - """The runtime this target is written for (it is not in the checkpoint): routing declines every other one - (``routing._runtime_mismatch``), so this is the backstop for a construction that did not go through routing. + """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. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_fused_decode_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_fused_decode_attn.py index 9a882ea6b933..1e1084aeccdf 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_fused_decode_attn.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_fused_decode_attn.py @@ -366,7 +366,7 @@ class M3FusedDecodeAttn(OpWrapper): the split-K LSE merge, writing the bf16 output and its MXFP8 copy for the o_proj GEMM. ``(tokens_per_cluster, blocks_per_cta)`` is the launch geometry; the target passes - ``kernels.m3_fused_attn.default_geometry(qlen, num_requests)``. Returns None: the cells read ``output``, + ``ops.m3_fused_attn.default_geometry(qlen, num_requests)``. Returns None: the cells read ``output``, ``mxfp8_data`` / ``mxfp8_scale`` and the three caches. """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_index_score.py index e7e57c10bde2..c8a9b54d2a3a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_index_score.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_index_score.py @@ -4,7 +4,7 @@ # 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 ``kernels/m3_index_score``.""" +the block's 128 index-K positions, with the pre-wait schedule of ``ops/m3_index_score``.""" from __future__ import annotations @@ -28,7 +28,7 @@ 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 ...kernels.m3_index_score import IndexScoreRunner + from ..m3_index_score import IndexScoreRunner if key not in IndexScoreRunner.kernel_cache: IndexScoreRunner._compile(*key) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode.py index 3c59454dac0a..b82d733c263b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/attention/m3_sparse_decode.py @@ -10,7 +10,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from ...kernels.triton_sparse_decode import minimax_m3_sparse_attn_decode +from ..triton_sparse_decode import minimax_m3_sparse_attn_decode from .m3_sparse_decode_lc import PAGE, sparse_attention_reference diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_mc_allreduce_norm.py index a2b74d268d09..0da06e59b14e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_mc_allreduce_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/comm/m3_mc_allreduce_norm.py @@ -128,7 +128,7 @@ class M3McAllreduceNorm(OpWrapper): spec=dict(m=48, maxt=48, opts=15, x2=True, quant=True, prepushed=0), ), # Every other row count of the decode boundaries (dispatch.MC_ALLREDUCE): the bucket and epilogue options the - # target picks for it (kernels/m3_mc_allreduce), the second addend and the MXFP8 copy alternating. + # 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", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding_firstnorm/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/embedding_firstnorm/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding_firstnorm/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/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_103__tp4/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding_firstnorm/op.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/embedding_firstnorm/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/embedding_firstnorm/op.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv.py index 5c1b2ada2d6d..a4cbb12d8085 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv.py @@ -10,8 +10,8 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from ...kernels.mx_gemv import op as _gemv 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"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv_bf16.py index 0951d082793a..1c6bbec4c10a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv_bf16.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/gemm/m3_mx_gemv_bf16.py @@ -10,8 +10,8 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from ...kernels.mx_gemv import op as _gemv 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, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/long_select.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/long_select.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/long_select.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_attn_lc.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_lc.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_attn_lc.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_attn_mega.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_attn_mega.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega_c1.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_attn_mega_c1.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_attn_mega_c1.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_attn_mega_c1.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_draft_step.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_draft_step.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_draft_step.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_force_vec.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_force_vec.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_force_vec.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_fused_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_fused_attn.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_fused_attn.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_fused_attn.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_index_score.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_index_score.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_index_score.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_mc_allreduce.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_mc_allreduce.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_mc_allreduce.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/m3_moe16_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/m3_moe16_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/m3_moe16_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/m3_moe16_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/m3_moe_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/m3_moe_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/m3_moe_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/op.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/op.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/op16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/op16.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_moe/op16.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_moe/op16.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_prequant.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_prequant.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_prequant.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_router_gemv.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_router_gemv.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_router_gemv.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_sampled_accept.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_sampled_accept.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_sampled_accept.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_shard_accept.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_shard_accept.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_shard_accept.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_spec_tail.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_spec_tail.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_spec_tail.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_target_head.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/m3_target_head.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_target_head.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe16_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe16_engine.py index 0e1595ed8ccf..950891839022 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe16_engine.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe16_engine.py @@ -12,8 +12,8 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from ...kernels.m3_moe.op import M3MoeLayer -from ...kernels.m3_moe.op16 import M3Moe16State +from ..m3_moe.op import M3MoeLayer +from ..m3_moe.op16 import M3Moe16State from .m3_moe_engine import moe_reference __all__ = ["M3Moe16State", "m3_moe16_engine"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe_engine.py index 219537a853f4..b6b0111b949f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe_engine.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/moe/m3_moe_engine.py @@ -15,7 +15,7 @@ get_shuffle_matrix_a_row_indices, ) -from ...kernels.m3_moe.op import ( +from ..m3_moe.op import ( HIDDEN, INTER, SHARED_INTER, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/msa_kv_update.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/msa_kv_update.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/msa_kv_update.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mx_gemv/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mx_gemv/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mx_gemv/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mx_gemv/m3_mx_gemv_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mx_gemv/m3_mx_gemv_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mx_gemv/m3_mx_gemv_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mx_gemv/op.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mx_gemv/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mx_gemv/op.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mxfp8_low_m_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mxfp8_low_m_gemm.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/mxfp8_low_m_gemm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/mxfp8_low_m_gemm.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/m3_embedding_firstnorm.py index 535f928018f5..1842171bde1a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/m3_embedding_firstnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/norm/m3_embedding_firstnorm.py @@ -10,7 +10,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from ...kernels.embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm +from ..embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm __all__ = ["DecodeEmbedding", "m3_embedding_firstnorm"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_force_vec.py index 1905516a8403..af4b06e3999f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_force_vec.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/spec/m3_force_vec.py @@ -8,7 +8,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from ...kernels.m3_force_vec import force_vec_kernel +from ..m3_force_vec import force_vec_kernel class M3ForceVec(OpWrapper): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/triton_sparse_decode.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/triton_sparse_decode.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/triton_sparse_decode.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/triton_sparse_decode_mq.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/kernels/triton_sparse_decode_mq.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/triton_sparse_decode_mq.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/__init__.py deleted file mode 100644 index 093ae922e88a..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Kernels of the MiniMax-M3 NVFP4 / sm_103 / tp8 target.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py index e2eedd3a8a17..c6a1b428c01f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py @@ -42,7 +42,7 @@ **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 for the -o_proj. It publishes it in ``kernels/m3_prequant.SLOT`` keyed by the BF16 tensor; the next +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. @@ -147,18 +147,13 @@ from tensorrt_llm.quantization.mode import QuantAlgo from . import weights as _weights -from .kernels import m3_draft_step as m3_draft_step_kernel -from .kernels import m3_fused_score_select as _fused_score_select_kernel -from .kernels import m3_index_score as _index_score_kernel -from .kernels import m3_mc_allreduce, m3_prequant, m3_spec_tail -from .kernels import m3_sampled_accept as m3_sampled_accept_kernel -from .kernels import m3_shard_accept as m3_shard_accept_kernel -from .kernels import msa_kv_update as _msa_kv_update_kernel -from .kernels.long_select import long_select_blocks -from .kernels.m3_router_gemv import m3_router_gemv as _router_gemv -from .kernels.m3_router_gemv import m3_router_gemv_fp4 as _router_gemv_fp4 -from .kernels.m3_router_gemv import m3_router_tc as _router_tc -from .kernels.mx_gemv import op as mx_gemv +from .ops import m3_draft_step as m3_draft_step_kernel +from .ops import m3_fused_score_select as _fused_score_select_kernel +from .ops import m3_index_score as _index_score_kernel +from .ops import m3_mc_allreduce, m3_prequant, m3_spec_tail +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.activation.m3_silu_and_mul import m3_silu_and_mul from .ops.attention.m3_dense_decode import m3_dense_decode from .ops.attention.m3_draft_attention import m3_draft_attention @@ -190,10 +185,15 @@ 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.m3_router_gemv import m3_router_gemv as _router_gemv +from .ops.m3_router_gemv import m3_router_gemv_fp4 as _router_gemv_fp4 +from .ops.m3_router_gemv import m3_router_tc as _router_tc from .ops.moe.m3_moe_engine import M3MoeState, ar_push_desc, m3_moe_engine from .ops.moe.m3_router_gemv import m3_router_gemv from .ops.moe.m3_router_gemv_fp4 import m3_router_gemv_fp4 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_step import MAX_TOKENS as DRAFT_STEP_KERNEL_MAX_TOKENS @@ -226,7 +226,7 @@ class DecodePlan: # 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 (kernels/m3_mc_allreduce.BUCKETS). +# verify; verify: the c4 / c8 verifies. Each is one multicast all-reduce bucket (ops/m3_mc_allreduce.BUCKETS). DECODE_PLANS = ( DecodePlan("tiny", 8, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), DecodePlan("small", 16, attn_mc_ar=True, moe_mc_ar=True, qkv_prequant=True), @@ -265,7 +265,7 @@ class DecodePlan: ROUTER_GEMV_MAX_TOKENS = 1 # capacity: fixed: measured router crossover ROUTER_TC_MAX_TOKENS = 48 # capacity: fixed: the tensor-core router's built rows # 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 (kernels/m3_shard_accept), at the configured draft length. The lm_head +# over the target lm_head's BF16 shard logits (ops/m3_shard_accept), at the configured draft length. The lm_head # is the in-tree BF16 GEMM: at the TP8 shard (25008 rows) a coded copy of the head saves no time. SHARD_ACCEPT_DRAFT_LEN = 3 # capacity: fixed: the built Eagle3 draft length SHARD_ACCEPT_MAX_REQUESTS = DECODE_MAX_TOKENS // ( @@ -276,14 +276,14 @@ class DecodePlan: # its GEMM lands beside the target tail it was meant to hide under. FC_EARLY_MAX_ROWS = 0 # capacity: fixed: measured fork crossover (never at TP8) # 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 (kernels/m3_spec_tail.draft_input), the drafter layer's attention, MLP and +# the draft input in one launch (ops/m3_spec_tail.draft_input), the drafter layer's attention, MLP and # all-reduces, and one kernel for the BF16 draft head, the cross-rank argmax and the pass's bookkeeping -# (kernels/m3_draft_step), at the configured draft length. The head kernel takes one row per request, at most 8, and +# (ops/m3_draft_step), at the configured draft length. The head kernel takes one row per request, at most 8, and # streams the head's BF16 weight fragments. DRAFT_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN # capacity: fixed: the built draft length DRAFT_STEP_MAX_REQUESTS = DRAFT_STEP_KERNEL_MAX_TOKENS # The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows -# (kernels/triton_sparse_decode.MQ_MIN_TOKENS). +# (ops/triton_sparse_decode.MQ_MIN_TOKENS). SPARSE_DECODE_MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query crossover # Decode spans of at most this many requests score and select their top-16 blocks in one launch # (ops/attention/m3_fused_score_select), which loads the older index-K pages before its grid wait; larger spans run @@ -759,8 +759,8 @@ def _text_model_config(model_config: ModelConfig) -> ModelConfig: def _validate_runtime(model_config: ModelConfig) -> None: - """The runtime this target is written for (it is not in the checkpoint): routing declines every other one - (``routing._runtime_mismatch``), so this is the backstop for a construction that did not go through routing. + """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. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_index_score.py index 7c9850fb3310..f5c385977d4c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_index_score.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_index_score.py @@ -4,7 +4,7 @@ # 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 ``kernels/m3_index_score``.""" +the block's 128 index-K positions, with the pre-wait schedule of ``ops/m3_index_score``.""" from __future__ import annotations @@ -28,7 +28,7 @@ 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 ...kernels.m3_index_score import IndexScoreRunner + from ..m3_index_score import IndexScoreRunner if key not in IndexScoreRunner.kernel_cache: IndexScoreRunner._compile(*key) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode.py index 648c03e76089..0438d20109b5 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/attention/m3_sparse_decode.py @@ -10,7 +10,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from ...kernels.triton_sparse_decode import minimax_m3_sparse_attn_decode +from ..triton_sparse_decode import minimax_m3_sparse_attn_decode from .m3_sparse_decode_lc8 import PAGE, sparse_attention_reference diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_mc_allreduce_norm.py index 2aad4aa72415..f7054de397dc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_mc_allreduce_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/comm/m3_mc_allreduce_norm.py @@ -128,7 +128,7 @@ class M3McAllreduceNorm(OpWrapper): spec=dict(m=48, maxt=48, opts=15, x2=True, quant=True, prepushed=0), ), # Every other row count of the decode boundaries (dispatch.MC_ALLREDUCE): the bucket and epilogue options the - # target picks for it (kernels/m3_mc_allreduce), the second addend and the MXFP8 copy alternating. + # 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", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding_firstnorm/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/embedding_firstnorm/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding_firstnorm/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/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_103__tp8/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding_firstnorm/op.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/embedding_firstnorm/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/embedding_firstnorm/op.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv.py index 0aad4c8ef78b..17c48fe4f7c6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv.py @@ -10,8 +10,8 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from ...kernels.mx_gemv import op as _gemv 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"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv_bf16.py index bbb0a98521d1..08abc2249da3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv_bf16.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/gemm/m3_mx_gemv_bf16.py @@ -10,8 +10,8 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from ...kernels.mx_gemv import op as _gemv 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, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/long_select.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/long_select.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/long_select.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_attn_lc8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_attn_lc8.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_attn_lc8.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_attn_lc8.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_draft_step.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_draft_step.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_draft_step.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_force_vec.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_force_vec.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_force_vec.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_fused_score_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_fused_score_select.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_fused_score_select.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_fused_score_select.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_index_score.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_index_score.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_index_score.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_mc_allreduce.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_mc_allreduce.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_mc_allreduce.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_moe/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_moe/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_moe/m3_moe_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/m3_moe_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_moe/m3_moe_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_moe/op.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_moe/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_moe/op.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_prequant.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_prequant.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_prequant.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_router_gemv.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_router_gemv.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_router_gemv.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_sampled_accept.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_sampled_accept.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_sampled_accept.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_shard_accept.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_shard_accept.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_shard_accept.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_spec_tail.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/m3_spec_tail.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_spec_tail.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_moe_engine.py index f46935ffe8ca..37d5ab5c6c46 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_moe_engine.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/moe/m3_moe_engine.py @@ -16,7 +16,7 @@ get_shuffle_matrix_a_row_indices, ) -from ...kernels.m3_moe.op import ( +from ..m3_moe.op import ( HIDDEN, INTER, SHARED_INTER, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/msa_kv_update.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/msa_kv_update.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/msa_kv_update.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/mx_gemv/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/mx_gemv/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/mx_gemv/m3_mx_gemv_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/m3_mx_gemv_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/mx_gemv/m3_mx_gemv_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/mx_gemv/op.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/mx_gemv/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/mx_gemv/op.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/m3_embedding_firstnorm.py index 535f928018f5..1842171bde1a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/m3_embedding_firstnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/norm/m3_embedding_firstnorm.py @@ -10,7 +10,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from ...kernels.embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm +from ..embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm __all__ = ["DecodeEmbedding", "m3_embedding_firstnorm"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_force_vec.py index 1905516a8403..af4b06e3999f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_force_vec.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/spec/m3_force_vec.py @@ -8,7 +8,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from ...kernels.m3_force_vec import force_vec_kernel +from ..m3_force_vec import force_vec_kernel class M3ForceVec(OpWrapper): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/triton_sparse_decode.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/triton_sparse_decode.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/triton_sparse_decode.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/triton_sparse_decode_mq.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/kernels/triton_sparse_decode_mq.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/triton_sparse_decode_mq.py 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 4d5623d7f834..867001663065 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 @@ -7,9 +7,11 @@ fits it raises ``NoTarget`` stating those facts. In ``auto`` the engine then uses the built-in MiniMax-M3 implementation; in ``require`` it raises that error. -The sparse-attention runtime is part of the decision too (``ROUTES_ON_RUNTIME``): -the targets are written for the MSA implementation with an FP8 indexer and an -FP8 or NVFP4 main KV cache, and any other runtime is the built-in model's. +The KV cache the engine runs with is part of the decision: the targets are +written for an FP8 or NVFP4 main KV cache. The rest of their runtime (the MSA +sparse attention implementation with an FP8 indexer) is an engine argument +rather than a checkpoint fact, so the targets check it at construction +(``_validate_runtime``). The checkpoint declares the vision-language wrapper, and the language model's shape lives in its ``text_config``, so that is what this tree reads. The @@ -18,17 +20,12 @@ from __future__ import annotations -import json -from pathlib import Path from typing import Any, Optional, Tuple from tensorrt_llm.llmapi.llm_args import EagleDecodingConfig from tensorrt_llm.quantization.mode import QuantAlgo -from ..._router_index import NoTarget, RuntimeFacts - -# ``route`` also reads the engine's sparse-attention runtime and KV cache dtype (``RuntimeFacts``). -ROUTES_ON_RUNTIME = True +from ..._router_index import NoTarget # The one GPU architecture these targets are written for: the decode kernels # (the tcgen05 GEMVs, the fused MoE engines, the attention megakernels and the @@ -43,16 +40,13 @@ # counts fix the per-rank operand shapes the decode kernels are compiled for. _SHAPE = (60, 6144, 128, 3072, 64, 4) -# The Eagle3 drafter geometry the targets' drafter kernels and weight tables are built for (the GQA drafter: the -# target's hidden size, heads, KV heads and vocabulary, a 18432-wide MLP). -_DRAFTER_GEOMETRY = dict( - hidden_size=6144, - num_attention_heads=64, - num_key_value_heads=4, - head_dim=128, - intermediate_size=18432, - vocab_size=200064, -) +# The architecture class and the shape do not pin everything the targets are built for: a newer +# checkpoint revision of this model family has the same architecture class and shapes but a different +# quant packing (compressed-tensors NVFP4 experts). The targets build a 5e6 RoPE table and dense +# attention on layers 0-2 (sparse on 3-59), so these two facts, with the quantization below, keep any +# other revision off the targets. +_ROPE_THETA = 5_000_000 +_SPARSE_ATTENTION_FREQ = (0,) * 3 + (1,) * 57 def _field(config: Any, name: str) -> Any: @@ -62,6 +56,28 @@ def _field(config: Any, name: str) -> Any: return getattr(config, name, None) +def _layout_mismatch(text_config: Any) -> Optional[str]: + """Why the text model's RoPE base or sparse-attention layers are not the ones the targets build, or None.""" + if text_config is None: + return None # the shape check declines it + theta = _field(text_config, "rope_theta") + if theta is None: + rope = _field(text_config, "rope_parameters") or {} + theta = _field( + rope.get("full_attention", rope) if isinstance(rope, dict) else rope, "rope_theta" + ) + sparse = _field(text_config, "sparse_attention_config") + freq = None if sparse is None else _field(sparse, "sparse_attention_freq") + # The built-in model reads use_sparse_attention=False as dense attention on every layer. + enabled = sparse is not None and _field(sparse, "use_sparse_attention") is not False + reasons = [] + if theta != _ROPE_THETA: + reasons.append(f"rope_theta {theta}, not {_ROPE_THETA}") + if not enabled or freq is None or tuple(int(f) for f in freq) != _SPARSE_ATTENTION_FREQ: + reasons.append("sparse attention layers other than 3-59") + return "; ".join(reasons) if reasons else None + + def _quant_algo(pretrained_config: Any) -> Optional[QuantAlgo]: """The quantization the checkpoint's config declares, as the engine reads it. @@ -114,107 +130,17 @@ def _eagle3_one_model_or_none(spec_config: Optional[Any]) -> bool: ) -def _drafter_mismatch(spec_config: Optional[Any]) -> Optional[str]: - """Why the speculative drafter is not the one the targets' flat drafter builds, or None when it is (or there is - no speculation). - - The targets draft with their own one-layer Eagle3 drafter (``M3Eagle3Draft``): one Llama-style layer, per-layer - fc_norms over several captured target layers, no eh_proj before attention or norm before the fc, no MLA, the - attention and MLP geometry of ``_DRAFTER_GEOMETRY``, the target's hidden size and vocabulary, no biases and an - untied head. Any other Eagle3 drafter is the built-in - implementation's. Read from the drafter checkpoint's ``config.json`` (the engine has resolved - ``speculative_model`` to a local directory before it builds the model). - """ - if spec_config is None: - return None - path = getattr(spec_config, "speculative_model", None) - if path is None: - return "no Eagle3 drafter checkpoint (speculative_model)" - try: - with open(Path(path) / "config.json") as f: - cfg = json.load(f) - except (OSError, ValueError) as err: - return f"the Eagle3 drafter's config.json is not readable ({err})" - eagle = cfg.get("eagle_config") or {} - hidden = cfg.get("hidden_size") - vocab = cfg.get("vocab_size") - captured = getattr(spec_config, "eagle3_layers_to_capture", None) - checks = ( - ( - cfg.get("num_hidden_layers") == 1, - f"{cfg.get('num_hidden_layers')} drafter layers, not 1", - ), - (bool(cfg.get("fc_norm", False)), "no per-layer fc_norm"), - (not eagle.get("eh_proj_before_attn", False), "eh_proj before attention"), - (not eagle.get("norm_before_fc", False), "norm before the fc"), - (not cfg.get("kv_lora_rank"), "an MLA drafter layer"), - ( - cfg.get("target_hidden_size", hidden) == hidden, - "a target hidden size other than the drafter's", - ), - ( - cfg.get("draft_vocab_size", vocab) in (None, vocab), - "a draft vocabulary other than the target's", - ), - ( - not (cfg.get("attention_bias") or cfg.get("mlp_bias") or cfg.get("bias")), - "biased projections", - ), - (not cfg.get("tie_word_embeddings", False), "a tied head"), - (captured is None or len(captured) > 1, "a single captured target layer"), - ) - heads = cfg.get("num_attention_heads") - defaults = dict( - num_key_value_heads=heads, - head_dim=hidden // heads if hidden and heads else None, - ) - geometry = [] - for name, want in _DRAFTER_GEOMETRY.items(): - value = cfg.get(name) - if value is None: - value = defaults.get(name) - geometry.append((value == want, f"{name}={value}, not {want}")) - reasons = [why for ok, why in (*checks, *geometry) if not ok] - return "; ".join(reasons) if reasons else None - - -def _runtime_mismatch(runtime: RuntimeFacts) -> Optional[str]: - """Why the engine's runtime is not the one the targets are written for, or None when it is. - - The MiniMax-M3 MSA backend (it owns the step metadata and the index-K cache), an FP8 indexer and an FP8 or - NVFP4 main KV cache, with the fused five-way QKV/index projection or the separate projections. Every other - runtime the engine accepts (the Triton implementation, the BF16 indexer, the BF16 main KV cache) is the built-in - implementation's. The targets' ``_validate_runtime`` checks the same at construction. - """ - sparse = runtime.sparse_attention_config - if sparse is None or _field(sparse, "algorithm") != "minimax_m3": - return "no MiniMax-M3 sparse_attention_config" - reasons = [] - implementation = _field(sparse, "implementation") - if implementation != "msa": - reasons.append(f"sparse attention implementation {implementation!r}, not 'msa'") - indexer = _field(sparse, "indexer_kv_dtype") - if indexer != "fp8": - reasons.append(f"indexer_kv_dtype {indexer!r}, not 'fp8'") - kv = runtime.kv_cache_quant_algo - if kv not in (QuantAlgo.FP8, QuantAlgo.NVFP4): - reasons.append(f"a {kv if kv is not None else 'BF16'} main KV cache, not FP8 or NVFP4") - return "; ".join(reasons) if reasons else None - - def route( pretrained_config: Any, mapping: Any, compute_capability: Tuple[int, int], spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], - *, - runtime: RuntimeFacts, ) -> type: """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. - ``kv_cache_quant_algo`` is the resolver's main KV-cache algorithm, passed to every routing module; this router - reads the same value, with the sparse attention config, from ``runtime`` (``ROUTES_ON_RUNTIME``). + ``kv_cache_quant_algo`` is the KV-cache algorithm the engine will run with + (see ``_router_index._kv_cache_quant_algo``), not the checkpoint's alone. Each target is imported where it is returned, so declining costs no import of modeling code. @@ -237,20 +163,20 @@ def route( ) quant_algo = _quant_algo(pretrained_config) parallel = _parallel(mapping) + layout_miss = _layout_mismatch(text_config) speculation = _eagle3_one_model_or_none(spec_config) - drafter = _drafter_mismatch(spec_config) if speculation else None - runtime_miss = _runtime_mismatch(runtime) + kv_cache = kv_cache_quant_algo in (QuantAlgo.FP8, QuantAlgo.NVFP4) # MiniMax-M3 NVFP4 on 4 GB300s (sm103), tensor-parallel attention and # experts, one-model Eagle3 or no speculative decoding. if ( compute_capability == _SM and shape == _SHAPE + and layout_miss is None and quant_algo == QuantAlgo.MIXED_PRECISION and parallel == "tp4" and speculation - and drafter is None - and runtime_miss is None + and kv_cache ): from .minimax_m3_nvfp4__sm_103__tp4.modeling import ModelingV2MiniMaxM3Nvfp4Sm103Tp4 @@ -261,27 +187,26 @@ def route( if ( compute_capability == _SM and shape == _SHAPE + and layout_miss is None and quant_algo == QuantAlgo.MIXED_PRECISION and parallel == "tp8" and speculation - and drafter is None - and runtime_miss is None + and kv_cache ): from .minimax_m3_nvfp4__sm_103__tp8.modeling import ModelingV2MiniMaxM3Nvfp4Sm103Tp8 return ModelingV2MiniMaxM3Nvfp4Sm103Tp8 raise NoTarget( - f"compute capability {compute_capability}, text_config shape {shape}, " + f"compute capability {compute_capability}, text_config shape {shape}" + f"{f' ({layout_miss})' if layout_miss else ''}, " f"quantization {quant_algo.name if quant_algo is not None else 'none'}, parallel " f"ws={mapping.world_size} tp={mapping.tp_size} pp={mapping.pp_size} cp={mapping.cp_size} " f"moe_tp={mapping.moe_tp_size} moe_ep={mapping.moe_ep_size} " 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"{f' (Eagle3 drafter: {drafter})' if drafter else ''}" - f"{f', runtime: {runtime_miss}' if runtime_miss else ''}; the " + 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 (a one-layer drafter with per-layer fc_norms) " - "or no speculative decoding, on the MSA sparse attention runtime with an FP8 indexer and an FP8 or NVFP4 KV " - "cache" + "tensor-parallel attention and experts, with one-model Eagle3 or no speculative decoding, and an FP8 or " + "NVFP4 KV cache" ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_c1.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_c1.py index 61f59e77ced0..202a2e2a8543 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_c1.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_c1.py @@ -8,9 +8,6 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels.m3_attn_mega_c1 import ( # noqa: E501 - MegaWorkspace, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_attn_c1 import ( # noqa: E501 m3_attn_c1 as op, ) @@ -19,6 +16,9 @@ mxfp8_rows, unswizzle_scales, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.m3_attn_mega_c1 import ( # noqa: E501 + MegaWorkspace, +) __extra_import_path__ = ["../../../..", "."] from _m3_fused_attn_step import ( # noqa: E402 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_mega.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_mega.py index b5a48b0252eb..d470e2c52847 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_mega.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_attn_mega.py @@ -8,9 +8,6 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels.m3_attn_mega import ( # noqa: E501 - MegaWorkspace, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_attn_mega import ( # noqa: E501 m3_attn_mega as op, ) @@ -20,6 +17,9 @@ reference_attention, unswizzle_scales, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.m3_attn_mega import ( # noqa: E501 + MegaWorkspace, +) __extra_import_path__ = ["../../../..", "."] from _m3_fused_attn_step import ( # noqa: E402 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_fused_decode_attn.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_fused_decode_attn.py index 4db2720f9564..f5fd205acd48 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_fused_decode_attn.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_m3_fused_decode_attn.py @@ -8,9 +8,6 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels.m3_fused_attn import ( # noqa: E501 - default_geometry, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_fused_decode_attn import ( # noqa: E501 PAGE, mxfp8_rows, @@ -19,6 +16,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_fused_decode_attn import ( # noqa: E501 m3_fused_decode_attn as op, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.m3_fused_attn import ( # noqa: E501 + default_geometry, +) __extra_import_path__ = ["../../../..", "."] from _m3_fused_attn_step import ( # noqa: E402 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_msa_kv_update.py index c5a2ec3f5af6..e1146b79a8b3 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_msa_kv_update.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/attention/test_modeling_v2_msa_kv_update.py @@ -7,15 +7,15 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels.msa_kv_update import ( # noqa: E501 - msa_kv_update as hook, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__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_103__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 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py index 1b749e0c7482..1b3888c2a638 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -552,7 +552,7 @@ def _run_one_rank() -> int: 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_103__tp4.kernels.m3_mc_allreduce as _ar # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.m3_mc_allreduce as _ar # noqa: E501 from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.comm import ( # noqa: E501 m3_mc_allreduce_norm as _plain, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_target_head.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_target_head.py index dd55be9a924b..69608913f82e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_target_head.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/gemm/test_modeling_v2_m3_target_head.py @@ -15,7 +15,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__tp4.ops.gemm.m3_target_head import ( # noqa: E501 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_input.py index ca13b3efb063..7863cf2e33cb 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_input.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_input.py @@ -6,7 +6,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops import ( # noqa: E501 m3_spec_tail, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_draft_input import ( # noqa: E501 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_step.py index 891889f3cd5a..e2a419a1d3a1 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_step.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_draft_step.py @@ -18,7 +18,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__tp4.ops.gemm.m3_target_head import ( # noqa: E501 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py index 249679ed08fe..d88cea0ad9f5 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/spec/test_modeling_v2_m3_exact_numerics.py @@ -23,7 +23,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops import ( # noqa: E501 m3_spec_tail, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 @@ -285,5 +285,5 @@ def test_the_drafter_is_the_in_tree_drafter() -> None: ): assert not hasattr(module, name), f"{module.__name__}.{name}" pkg = modeling.__name__.rsplit(".", 1)[0] - for gone in ("kernels.m3_draft_head", "ops.spec.m3_draft_head", "ops.gemm.m3_draft_gemv"): + 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_103__tp4/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_dispatch.py index 1fa890f5f127..450acab7c7d4 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_dispatch.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_dispatch.py @@ -17,7 +17,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4 import ( # noqa: E501 modeling as dispatch, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops import ( # noqa: E501 m3_attn_mega_c1, m3_draft_step, m3_fused_attn, @@ -26,12 +26,9 @@ m3_target_head, triton_sparse_decode, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops import ( # noqa: E501 mxfp8_low_m_gemm as low_m, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels.mx_gemv import ( # noqa: E501 - op as mx_gemv, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.gemm import ( # noqa: E501 m3_target_head as target_head_entry, ) @@ -47,6 +44,9 @@ SHARED_DOWN, SHARED_GATE_UP, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.mx_gemv import ( # noqa: E501 + op as mx_gemv, +) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.spec import ( # noqa: E501 m3_draft_step as draft_step_entry, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_prebuild.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_prebuild.py index 1a547c0039ac..df4afe2f15b3 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_prebuild.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_prebuild.py @@ -17,12 +17,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__tp4.kernels.mx_gemv import ( # noqa: E501 - op as mx_gemv, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.attention.m3_index_score import ( # noqa: E501 m3_index_score, ) @@ -31,6 +28,9 @@ SHARED_DOWN, SHARED_GATE_UP, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops.mx_gemv import ( # noqa: E501 + op as mx_gemv, +) dispatch = modeling diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_rank_agreement.py index 170597182d6b..ec23eb22fc9d 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_rank_agreement.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/test_modeling_v2_m3_rank_agreement.py @@ -14,7 +14,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp4.ops import ( # noqa: E501 m3_mc_allreduce as mc, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_fused_score_select.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_fused_score_select.py index f8e2bf4bfacf..ea21f436a343 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_fused_score_select.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_m3_fused_score_select.py @@ -18,7 +18,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 modeling as dispatch, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops import ( # noqa: E501 m3_fused_score_select as fss, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_fused_score_select import ( # noqa: E501 @@ -157,7 +157,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res @pytest.mark.parametrize("cell", op.CELLS[:2], ids=[c.why[:44] for c in op.CELLS[:2]]) def test_the_target_wrapper_is_the_entry(cell) -> None: - """``kernels/m3_fused_score_select``: the geometry check admits the cell, the scratch and counters are created + """``ops/m3_fused_score_select``: the geometry check admits the cell, the scratch and counters are created once per (device, engine block count), and the wrapper's table is the entry's.""" spec = cell.spec inputs = _inputs(spec, 3) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_msa_kv_update.py index ffca6d60eefd..458a7000d9fd 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_msa_kv_update.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/attention/test_modeling_v2_msa_kv_update.py @@ -7,15 +7,15 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels.msa_kv_update import ( # noqa: E501 - msa_kv_update as hook, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__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_103__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 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py index d322305df03b..8a77ae91b481 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -559,7 +559,7 @@ def _run_one_rank() -> int: 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_103__tp8.kernels.m3_mc_allreduce as _ar # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.m3_mc_allreduce as _ar # noqa: E501 from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.comm import ( # noqa: E501 m3_mc_allreduce_norm as _plain, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_input.py index 794deaf1678d..7476e6f10e5e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_input.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_input.py @@ -6,7 +6,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops import ( # noqa: E501 m3_spec_tail, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_draft_input import ( # noqa: E501 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_step.py index 339b61ce3dc0..6cf9240e71a0 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_step.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_draft_step.py @@ -17,7 +17,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__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_103__tp8.ops.spec.m3_draft_step import ( # noqa: E501 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_exact_numerics.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_exact_numerics.py index f705462595cc..af6e3c862fd7 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_exact_numerics.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/spec/test_modeling_v2_m3_exact_numerics.py @@ -25,7 +25,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops import ( # noqa: E501 m3_spec_tail, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 @@ -289,11 +289,11 @@ def test_the_drafter_is_the_in_tree_drafter() -> None: assert not hasattr(modeling.M3Eagle3Worker, "prebuild_target_head") pkg = modeling.__name__.rsplit(".", 1)[0] for gone in ( - "kernels.m3_draft_head", + "ops.m3_draft_head", "ops.spec.m3_draft_head", "ops.gemm.m3_draft_gemv", - "kernels.m3_target_head", + "ops.m3_target_head", "ops.gemm.m3_target_head", - "kernels.m3_draft_lowm", + "ops.m3_draft_lowm", ): 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_103__tp8/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_dispatch.py index 2c098e15ee5b..9f7508bfdd59 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_dispatch.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_dispatch.py @@ -18,7 +18,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 modeling as dispatch, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops import ( # noqa: E501 m3_attn_lc8, m3_draft_step, m3_fused_score_select, @@ -26,9 +26,6 @@ m3_shard_accept, triton_sparse_decode, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels.mx_gemv import ( # noqa: E501 - op as mx_gemv, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention import ( # noqa: E501 m3_fused_score_select as fused_score_select_entry, ) @@ -45,6 +42,9 @@ QKV, QKV_SPARSE, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.mx_gemv import ( # noqa: E501 + op as mx_gemv, +) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.spec import ( # noqa: E501 m3_draft_step as draft_step_entry, ) @@ -135,12 +135,12 @@ def test_the_removed_tp4_paths_are_gone() -> None: ): assert not hasattr(dispatch, constant), constant for module in ( - "kernels.m3_target_head", - "kernels.mxfp8_low_m_gemm", - "kernels.m3_attn_mega", - "kernels.m3_attn_mega_c1", - "kernels.m3_fused_attn", - "kernels.m3_attn_lc", + "ops.m3_target_head", + "ops.mxfp8_low_m_gemm", + "ops.m3_attn_mega", + "ops.m3_attn_mega_c1", + "ops.m3_fused_attn", + "ops.m3_attn_lc", "ops.gemm.m3_target_head", "ops.gemm.mxfp8_low_m_gemm", "ops.moe.m3_moe16_engine", diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_prebuild.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_prebuild.py index e932865571e5..d1ee4df80e69 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_prebuild.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_prebuild.py @@ -18,12 +18,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.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_103__tp8.kernels.mx_gemv import ( # noqa: E501 - op as mx_gemv, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.attention.m3_index_score import ( # noqa: E501 m3_index_score, ) @@ -32,6 +29,9 @@ O_PROJ, QKV_SPARSE, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops.mx_gemv import ( # noqa: E501 + op as mx_gemv, +) dispatch = modeling diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py index ee3c95ac63be..81972c1c503e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py @@ -14,7 +14,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_103__tp8.ops import ( # noqa: E501 m3_mc_allreduce as mc, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py index 8bc830b733ef..6d563ff9f831 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py @@ -29,20 +29,20 @@ #: (``ops/attention/m3_msa_attention.py`` left the list: its reference dequantizes NVFP4 pages with each target's own #: long-context sparse decode entry, ``m3_sparse_decode_lc`` at TP4 and ``m3_sparse_decode_lc8`` at TP8.) IDENTICAL_TO_TP4 = ( - "kernels/embedding_firstnorm/__init__.py", - "kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py", - "kernels/embedding_firstnorm/op.py", - "kernels/long_select.py", - "kernels/m3_force_vec.py", - "kernels/m3_index_score.py", - "kernels/m3_prequant.py", - "kernels/m3_router_gemv.py", - "kernels/m3_sampled_accept.py", - "kernels/m3_shard_accept.py", - "kernels/m3_spec_tail.py", - "kernels/msa_kv_update.py", - "kernels/mx_gemv/m3_mx_gemv_kernel.py", - "kernels/triton_sparse_decode_mq.py", + "ops/embedding_firstnorm/__init__.py", + "ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py", + "ops/embedding_firstnorm/op.py", + "ops/long_select.py", + "ops/m3_force_vec.py", + "ops/m3_index_score.py", + "ops/m3_prequant.py", + "ops/m3_router_gemv.py", + "ops/m3_sampled_accept.py", + "ops/m3_shard_accept.py", + "ops/m3_spec_tail.py", + "ops/msa_kv_update.py", + "ops/mx_gemv/m3_mx_gemv_kernel.py", + "ops/triton_sparse_decode_mq.py", "ops/__init__.py", "ops/activation/__init__.py", "ops/attention/__init__.py", 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 1f0a0ea6c085..d63924daf06b 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -564,123 +564,113 @@ def test_minimax_m3_parallel_names_the_target_topology(mapping_kwargs, parallel) assert routing._parallel(Mapping(**mapping_kwargs)) == parallel -# The MiniMax-M3 Eagle3 drafter of the released checkpoint (inferact MiniMax-M3-EAGLE3-GQA), as its config.json reads. -_M3_DRAFTER = dict( - num_hidden_layers=1, - fc_norm=True, - hidden_size=6144, - num_attention_heads=64, - num_key_value_heads=4, - head_dim=128, - intermediate_size=18432, - vocab_size=200064, - draft_vocab_size=200064, - attention_bias=False, - tie_word_embeddings=False, +# The MiniMax-M3 text model's shape fields, in ``routing._SHAPE`` order. +_M3_SHAPE_FIELDS = ( + "num_hidden_layers", + "hidden_size", + "num_local_experts", + "intermediate_size", + "num_attention_heads", + "num_key_value_heads", ) +# The MiniMax-M3 RoPE base and sparse-attention layers (dense 0-2, sparse 3-59), as its config.json reads. +_M3_LAYOUT = dict( + rope_theta=5_000_000, sparse_attention_config=dict(sparse_attention_freq=[0] * 3 + [1] * 57) +) +# Another checkpoint revision with the same architecture class and shape: another RoPE base and sparse layout. +_OTHER_LAYOUT = dict( + rope_theta=1_000_000, sparse_attention_config=dict(sparse_attention_freq=[0] * 4 + [1] * 56) +) + + +def _m3_text_config(**layout): + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + return dict(zip(_M3_SHAPE_FIELDS, routing._SHAPE), **(layout or _M3_LAYOUT)) -def _m3_drafter_spec(tmp_path, captured=None, **overrides): - import json + +@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): from types import SimpleNamespace - (tmp_path / "config.json").write_text(json.dumps({**_M3_DRAFTER, **overrides})) - return SimpleNamespace(speculative_model=str(tmp_path), eagle3_layers_to_capture=captured) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + monkeypatch.setattr(routing, "_eagle3_one_model_or_none", lambda spec: True) + spec_config = ( + None + if spec is None + else SimpleNamespace(spec_dec_mode=SimpleNamespace(name="EAGLE3_ONE_MODEL")) + ) + target = routing.route( + SimpleNamespace(text_config=_m3_text_config()), + Mapping(world_size=tp, tp_size=tp, moe_tp_size=tp, moe_ep_size=1), + routing._SM, + spec_config, + QuantAlgo.NVFP4, + ) + assert target.__name__ == f"ModelingV2MiniMaxM3Nvfp4Sm103Tp{tp}" @pytest.mark.parametrize( - "overrides, captured, reason", + "layout, reason", [ - ({}, None, None), - ({}, {1, 29, 56}, None), - (dict(num_hidden_layers=2), None, "2 drafter layers, not 1"), - (dict(fc_norm=False), None, "no per-layer fc_norm"), - (dict(eagle_config={"eh_proj_before_attn": True}), None, "eh_proj before attention"), - (dict(eagle_config={"norm_before_fc": True}), None, "norm before the fc"), - (dict(kv_lora_rank=512), None, "an MLA drafter layer"), - (dict(target_hidden_size=4096), None, "a target hidden size other than the drafter's"), - (dict(draft_vocab_size=32000), None, "a draft vocabulary other than the target's"), - (dict(mlp_bias=True), None, "biased projections"), - (dict(tie_word_embeddings=True), None, "a tied head"), - ({}, {-1}, "a single captured target layer"), - (dict(num_key_value_heads=1), None, "num_key_value_heads=1, not 4"), - (dict(intermediate_size=14336), None, "intermediate_size=14336, not 18432"), + (_OTHER_LAYOUT, "rope_theta 1000000, not 5000000; sparse attention layers other than 3-59"), + (dict(_M3_LAYOUT, rope_theta=1_000_000), "rope_theta 1000000, not 5000000"), ( - dict(hidden_size=4096, head_dim=None), - None, - "hidden_size=4096, not 6144; head_dim=64, not 128", + dict(_M3_LAYOUT, sparse_attention_config=_OTHER_LAYOUT["sparse_attention_config"]), + "sparse attention layers other than 3-59", + ), + ( + dict(sparse_attention_config=_M3_LAYOUT["sparse_attention_config"]), + "rope_theta None, not 5000000", + ), + ( + dict( + _M3_LAYOUT, + sparse_attention_config=dict( + _M3_LAYOUT["sparse_attention_config"], use_sparse_attention=False + ), + ), + "sparse attention layers other than 3-59", ), ], - ids=[ - "released_drafter", - "three_captured_layers", - "two_layers", - "no_fc_norm", - "eh_proj_before_attn", - "norm_before_fc", - "mla", - "target_hidden", - "draft_vocab", - "bias", - "tied_head", - "one_captured_layer", - "one_kv_head", - "mlp_width", - "hidden_size", - ], + ids=["other_revision", "rope_base", "sparse_layers", "no_rope_base", "sparse_disabled"], ) -def test_minimax_m3_routes_only_the_drafter_its_targets_build( - tmp_path, overrides, captured, reason -): - """The targets draft with their own one-layer fc_norm Eagle3 drafter: any other drafter is declined in routing - (``auto`` takes the built-in model, ``require`` names the criterion) instead of failing at construction.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing - - spec = _m3_drafter_spec(tmp_path, captured, **overrides) - assert routing._drafter_mismatch(spec) == reason - assert routing._drafter_mismatch(None) is None - - -@pytest.mark.parametrize("drafter_layers", [1, 3], ids=["flat_drafter", "multi_layer_drafter"]) -def test_minimax_m3_route_declines_a_drafter_it_does_not_build( - tmp_path, monkeypatch, drafter_layers -): +def test_minimax_m3_route_declines_another_rope_base_or_sparse_layout(monkeypatch, layout, reason): + """The architecture name and the shape do not pin the RoPE base or which layers are sparse: a checkpoint with + the target shape but another of either (another checkpoint revision) is not routed to the MiniMax-M3 targets, + even with their quantization.""" 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 - from tensorrt_llm.quantization.mode import QuantAlgo monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) - monkeypatch.setattr(routing, "_eagle3_one_model_or_none", lambda spec: True) - text_config = dict( - zip( - ( - "num_hidden_layers", - "hidden_size", - "num_local_experts", - "intermediate_size", - "num_attention_heads", - "num_key_value_heads", - ), - routing._SHAPE, + text_config = _m3_text_config(**layout) + assert routing._layout_mismatch(text_config) == reason + with pytest.raises(NoTarget, match=reason.split(";")[0]): + routing.route( + SimpleNamespace(text_config=text_config), + Mapping(world_size=4, tp_size=4, moe_tp_size=4, moe_ep_size=1), + routing._SM, + None, + QuantAlgo.NVFP4, ) - ) - spec = _m3_drafter_spec(tmp_path, num_hidden_layers=drafter_layers) - spec.spec_dec_mode = SimpleNamespace(name="EAGLE3_ONE_MODEL") - mapping = Mapping(world_size=4, tp_size=4, moe_tp_size=4, moe_ep_size=1) - pretrained = SimpleNamespace(text_config=text_config) - runtime = _m3_runtime() - if drafter_layers == 1: - target = routing.route( - pretrained, mapping, routing._SM, spec, runtime.kv_cache_quant_algo, runtime=runtime + + +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 + + sparse = _M3_LAYOUT["sparse_attention_config"] + assert routing._layout_mismatch(_m3_text_config()) is None + for rope in (dict(rope_theta=5_000_000), dict(full_attention=dict(rope_theta=5_000_000))): + assert ( + routing._layout_mismatch(dict(sparse_attention_config=sparse, rope_parameters=rope)) + is None ) - assert target.__name__ == "ModelingV2MiniMaxM3Nvfp4Sm103Tp4" - else: - with pytest.raises(NoTarget, match="Eagle3 drafter: 3 drafter layers, not 1"): - routing.route( - pretrained, mapping, routing._SM, spec, runtime.kv_cache_quant_algo, runtime=runtime - ) def _m3_sparse(**overrides): @@ -697,87 +687,12 @@ def _m3_sparse(**overrides): return MiniMaxM3SparseAttentionConfig(**fields) -def _m3_runtime(sparse=None, kv="NVFP4"): - from tensorrt_llm._torch._experimental.modeling_v2._router_index import RuntimeFacts - - return RuntimeFacts( - sparse_attention_config=_m3_sparse() if sparse is None else sparse, kv_cache_quant_algo=kv - ) - - -# Every runtime the engine accepts for MiniMax-M3 that the targets are not written for; the built-in model runs them. -_M3_RUNTIME_DECLINES = [ - ( - dict(implementation="triton", indexer_kv_dtype="bf16", fuse_qkv_index_projection=False), - "FP8", - "sparse attention implementation 'triton', not 'msa'; indexer_kv_dtype 'bf16', not 'fp8'", - ), - ( - dict(indexer_kv_dtype="bf16", fuse_qkv_index_projection=False), - "FP8", - "indexer_kv_dtype 'bf16', not 'fp8'", - ), - ( - dict(fuse_qkv_index_projection=False), - None, - "a BF16 main KV cache, not FP8 or NVFP4", - ), -] -_M3_RUNTIME_DECLINE_IDS = ["triton", "bf16_indexer", "bf16_main_kv"] - - -@pytest.mark.parametrize( - "sparse, kv", - [ - (dict(), "NVFP4"), - (dict(), "FP8"), - (dict(fuse_qkv_index_projection=False), "FP8"), - (dict(fuse_qkv_index_projection=False), "NVFP4"), - ], - ids=["fused_nvfp4", "fused_fp8", "unfused_fp8", "unfused_nvfp4"], -) -def test_minimax_m3_runtime_the_targets_are_written_for(sparse, kv): - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing - - assert routing._runtime_mismatch(_m3_runtime(_m3_sparse(**sparse), kv)) is None - - -@pytest.mark.parametrize("sparse, kv, reason", _M3_RUNTIME_DECLINES, ids=_M3_RUNTIME_DECLINE_IDS) -def test_minimax_m3_runtime_mismatch_names_the_setting(sparse, kv, reason): - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing - - assert routing._runtime_mismatch(_m3_runtime(_m3_sparse(**sparse), kv)) == reason - - -def test_minimax_m3_runtime_without_its_sparse_config_is_declined(): - from tensorrt_llm._torch._experimental.modeling_v2._router_index import RuntimeFacts - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing - - runtime = RuntimeFacts(sparse_attention_config=None, kv_cache_quant_algo="NVFP4") - assert routing._runtime_mismatch(runtime) == "no MiniMax-M3 sparse_attention_config" - - -def _m3_resolve_config(sparse, kv_checkpoint, tp=4): +def _m3_resolve_config(sparse, kv_checkpoint, tp=4, layout=None): """A ModelConfig as the loader holds it when it resolves: the MiniMax-M3 VL wrapper with the target text shape, the checkpoint's KV quantization (before kv_cache_config.dtype is applied) and the sparse-attention runtime.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing - from tensorrt_llm.models.modeling_utils import QuantConfig - - text_config = dict( - zip( - ( - "num_hidden_layers", - "hidden_size", - "num_local_experts", - "intermediate_size", - "num_attention_heads", - "num_key_value_heads", - ), - routing._SHAPE, - ) - ) pretrained = PretrainedConfig( - architectures=["MiniMaxM3SparseForConditionalGeneration"], text_config=text_config + architectures=["MiniMaxM3SparseForConditionalGeneration"], + text_config=_m3_text_config(**(layout or {})), ) return ModelConfig( pretrained_config=pretrained, @@ -804,7 +719,6 @@ def _minimax_m3_sm(_on_sm10x): def test_minimax_m3_pr_runtime_resolves_to_the_target(monkeypatch, mode, tp): """The PR deployment (MSA, FP8 indexer, fused projection, kv_cache_config.dtype='nvfp4') takes the target.""" from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing - from tensorrt_llm.quantization.mode import QuantAlgo monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) config = _m3_resolve_config(_m3_sparse(), None, tp) @@ -813,42 +727,52 @@ def test_minimax_m3_pr_runtime_resolves_to_the_target(monkeypatch, mode, tp): @pytest.mark.usefixtures("_minimax_m3_sm") -@pytest.mark.parametrize("sparse, kv, reason", _M3_RUNTIME_DECLINES, ids=_M3_RUNTIME_DECLINE_IDS) @pytest.mark.parametrize("tp", [4, 8]) -def test_minimax_m3_other_runtimes_fall_back_under_auto(monkeypatch, sparse, kv, reason, tp): - """``auto`` resolves nothing, so the loader builds the registered built-in MiniMax-M3 model.""" +def test_minimax_m3_bf16_kv_cache_falls_back_under_auto(monkeypatch, tp): + """The targets run an FP8 or NVFP4 KV cache: ``auto`` with the BF16 cache resolves nothing, so the loader builds + the registered built-in MiniMax-M3 model.""" from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing from tensorrt_llm._torch.models.modeling_minimaxm3 import MiniMaxM3VLForConditionalGeneration - from tensorrt_llm.quantization.mode import QuantAlgo monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) # kv_cache_config.dtype 'auto' keeps the checkpoint's KV quantization (None: the BF16 cache). - config = _m3_resolve_config(_m3_sparse(**sparse), kv, tp) + config = _m3_resolve_config(_m3_sparse(fuse_qkv_index_projection=False), None, tp) assert modeling_v2_resolve(config, _m3_llm_args("auto", "auto")) is None resolved = replace(config, modeling_v2_target=None) assert AutoModelForCausalLM._resolve_class(resolved) is MiniMaxM3VLForConditionalGeneration @pytest.mark.usefixtures("_minimax_m3_sm") -@pytest.mark.parametrize("sparse, kv, reason", _M3_RUNTIME_DECLINES, ids=_M3_RUNTIME_DECLINE_IDS) -def test_minimax_m3_other_runtimes_raise_under_require(monkeypatch, sparse, kv, reason): +def test_minimax_m3_bf16_kv_cache_raises_under_require(monkeypatch): from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing - from tensorrt_llm.quantization.mode import QuantAlgo monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) - config = _m3_resolve_config(_m3_sparse(**sparse), kv) + config = _m3_resolve_config(_m3_sparse(fuse_qkv_index_projection=False), None) with pytest.raises(ValueError) as excinfo: modeling_v2_resolve(config, _m3_llm_args("require", "auto")) message = str(excinfo.value) assert "MiniMaxM3SparseForConditionalGeneration" in message - assert f"runtime: {reason}" in message + assert "KV cache bf16" in message + + +@pytest.mark.usefixtures("_minimax_m3_sm") +@pytest.mark.parametrize("tp", [4, 8]) +def test_minimax_m3_other_checkpoint_revision_is_declined(monkeypatch, tp): + """Another checkpoint revision (same architecture class and shape, another RoPE base and sparse layout) with the + PR runtime: ``auto`` builds the built-in model and ``require`` names the layout.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + config = _m3_resolve_config(_m3_sparse(), None, tp, layout=_OTHER_LAYOUT) + assert modeling_v2_resolve(config, _m3_llm_args("auto", "nvfp4")) is None + with pytest.raises(ValueError, match="rope_theta 1000000, not 5000000"): + modeling_v2_resolve(config, _m3_llm_args("require", "nvfp4")) @pytest.mark.usefixtures("_minimax_m3_sm") def test_minimax_m3_kv_cache_dtype_overrides_the_checkpoint(monkeypatch): """An explicit kv_cache_config.dtype is the KV cache the engine runs with, whatever the checkpoint declares.""" from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing - from tensorrt_llm.quantization.mode import QuantAlgo monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) config = _m3_resolve_config(_m3_sparse(fuse_qkv_index_projection=False), None) From f71400b4b5f47a426640770a9d3ff055d9a1d8a0 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:33:41 -0700 Subject: [PATCH 12/14] [None][feat] modeling_v2 'default' mode: MiniMax-M3 GB300 targets on by default where measured The modeling_v2 argument gains a 'default' value, now its default. It behaves as 'auto' on the deployments the architecture's routing module enables through enabled_by_default(target, role), and as 'off' everywhere else. An architecture whose routing module has no enabled_by_default is not routed at all, so every other model is unchanged. MiniMax-M3 enables its GB300 TP4 / TP8 targets on the generation side of disaggregated serving. The role comes from the cache transceiver config and TRTLLM_DISAGG_ROLE. 'off' turns it off. route() stays the standard five-argument decision. The runtime checks the targets make at construction (_validate_runtime, the Eagle3 drafter) now raise NoTarget, a ValueError, and AutoModelForCausalLM.from_config builds the built-in implementation on such a refusal unless the target was resolved under 'require'. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../_experimental/modeling_v2/README.md | 3 +- .../modeling_v2/_router_index.py | 60 +++++++- .../minimax_m3_nvfp4__sm_103__tp4/modeling.py | 13 +- .../minimax_m3_nvfp4__sm_103__tp8/modeling.py | 13 +- .../modeling_v2/models/minimax_m3/routing.py | 25 +++- tensorrt_llm/_torch/models/modeling_auto.py | 29 +++- tensorrt_llm/llmapi/llm_args.py | 12 +- .../schemas/trtllm-serve-config.schema.json | 5 +- .../trtllm-serve-disagg-config.schema.json | 10 +- .../usage/llm_args_golden_manifest.json | 1 + .../modeling_v2/test_modeling_v2_routing.py | 129 +++++++++++++++++- .../api_stability/references/llm.yaml | 4 +- 12 files changed, 273 insertions(+), 31 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index 5e34934827bc..5801e6a58904 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md @@ -24,7 +24,8 @@ modeling_v2: require | `modeling_v2` | Behaviour | |---|---| -| `off` (default) | The resolver returns immediately. Nothing beyond it is imported and behaviour is byte-for-byte what it is today. | +| `default` (default) | `auto` where the family's routing module enables the target by default (`enabled_by_default(target, role)`, the serving roles it was measured in), `off` everywhere else. A family without `enabled_by_default` is not routed at all. | +| `off` | The resolver returns immediately. Nothing beyond it is imported and behaviour is byte-for-byte what it is today. | | `auto` | Uses a target when one claims this exact configuration; falls back to the built-in implementation otherwise, logging why. | | `require` | Raises instead of falling back, saying which criterion did not hold. | diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index 60a138b565f1..b7a2455c6861 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -87,6 +87,8 @@ def step_contract_enabled() -> bool: class ModelingV2Mode(str, enum.Enum): """What to do when a deployment reaches the modeling_v2 resolver.""" + #: ``auto`` on the deployments the routing module enables by default, ``off`` everywhere else. + DEFAULT = "default" OFF = "off" AUTO = "auto" REQUIRE = "require" @@ -96,15 +98,43 @@ def of(cls, llm_args: "TorchLlmArgs") -> "ModelingV2Mode": return cls(llm_args.modeling_v2) -class NoTarget(Exception): - """Raised by a routing module's ``route`` when no target fits. +class NoTarget(ValueError): + """Raised by a routing module's ``route`` when no target fits, or by a target at construction when the runtime + the engine arguments set is not the one it is written for. The message says, in words, which fact ruled the targets out. Under - ``auto`` the resolver logs it and falls back to the built-in implementation; + ``auto`` the resolver logs it and falls back to the built-in implementation + (for a refusal at construction, ``AutoModelForCausalLM.from_config`` does); under ``require`` it becomes the error the user sees. """ +class DeploymentRole(str, enum.Enum): + """What this engine serves: both phases, or one side of disaggregated serving.""" + + AGGREGATED = "aggregated" + CONTEXT = "context" + GENERATION = "generation" + #: Disaggregated, but the launcher did not say which side. + UNKNOWN = "unknown" + + +#: Set by ``trtllm-serve`` on disaggregated workers (``context`` or ``generation``; legacy ``ctx*`` / ``gen*``). +DISAGG_ROLE_ENV = "TRTLLM_DISAGG_ROLE" + + +def deployment_role(llm_args: "TorchLlmArgs") -> DeploymentRole: + """The role of this engine, known before the model is built: no cache transceiver means aggregated serving.""" + if llm_args.cache_transceiver_config is None: + return DeploymentRole.AGGREGATED + role = os.environ.get(DISAGG_ROLE_ENV, "").strip().lower() + if role == "context" or role.startswith("ctx"): + return DeploymentRole.CONTEXT + if role == "generation" or role.startswith("gen"): + return DeploymentRole.GENERATION + return DeploymentRole.UNKNOWN + + # The mode each target class was last resolved under (the loader builds the class right after resolving it). _RESOLVED_MODES: Dict[type, ModelingV2Mode] = {} @@ -187,6 +217,11 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti when ``llm_args.modeling_v2`` is ``off``, when no routing module claims the architecture, or when the routing module's ``route`` raises ``NoTarget``. In ``require`` every one of those but ``off`` raises ``ValueError`` instead. + + ``default`` resolves as ``auto`` and then keeps the target only where the + routing module's ``enabled_by_default(target, role)`` returns None (a reason + string declines it). An architecture whose routing module has no + ``enabled_by_default`` is not routed at all under ``default``. """ mode = ModelingV2Mode.of(llm_args) if mode is ModelingV2Mode.OFF: @@ -197,6 +232,11 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti return None arch = pretrained_config.architectures[0] + if mode is ModelingV2Mode.DEFAULT: + routing = routing_module(arch) + if getattr(routing, "enabled_by_default", None) is None: + return None + try: routing = routing_module(arch) if routing is None: @@ -218,6 +258,20 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti f"modeling_v2: no target fits {arch!r} ({why}); using the built-in implementation" ) return None + if target is not None and mode is ModelingV2Mode.DEFAULT: + role = deployment_role(llm_args) + why = routing.enabled_by_default(target, role) + if why is not None: + logger.info( + f"modeling_v2: {target.__name__} fits {arch!r} but is not enabled by default for this deployment " + f"({why}); using the built-in implementation. Set modeling_v2='auto' to use the target." + ) + return None + logger.info( + f"modeling_v2: {target.__name__} is enabled by default for {arch!r} ({role.value}); " + "set modeling_v2='off' to use the built-in implementation" + ) + mode = ModelingV2Mode.AUTO if target is not None: _RESOLVED_MODES[target] = mode return target diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py index 5fe483d48cc7..a7ac9f445c20 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py @@ -80,6 +80,7 @@ from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( ModelingV2Mode, + NoTarget, resolved_mode, step_contract_enabled, ) @@ -827,16 +828,16 @@ def _validate_runtime(model_config: ModelConfig) -> None: or getattr(sparse, "algorithm", None) != "minimax_m3" or sparse.implementation != "msa" ): - raise ValueError( + raise NoTarget( "this MiniMax-M3 target requires sparse_attention_config.algorithm='minimax_m3' with implementation='msa'" ) if getattr(sparse, "indexer_kv_dtype", None) != "fp8": - raise ValueError( + raise NoTarget( "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") + raise NoTarget("this MiniMax-M3 target requires an FP8 or NVFP4 KV cache") def _sparse_layer_ids(cfg: PretrainedConfig) -> List[int]: @@ -3384,7 +3385,7 @@ def __init__( or getattr(cfg, "bias", False) or getattr(cfg, "tie_word_embeddings", False) ): - raise ValueError( + raise NoTarget( "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" @@ -3430,7 +3431,7 @@ def __init__( sparse_params=None, ) if not self.attn.support_fused_rope() or not self.attn.support_fused_qkv(): - raise ValueError( + raise NoTarget( f"the drafter's attention backend {type(self.attn).__name__} must fuse RoPE and QKV" ) self.attn.update_quant_config(quant_config) @@ -4452,7 +4453,7 @@ def _build_speculation(shell, model_config: ModelConfig): or spec.eagle3_model_arch != "llama3" or getattr(spec, "use_dynamic_tree", False) ): - raise ValueError( + raise NoTarget( "this MiniMax-M3 target speculates with linear one-model Eagle3 (llama3 drafter) only" ) draft_config = ModelConfig.from_pretrained( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py index c6a1b428c01f..06887c682419 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py @@ -82,6 +82,7 @@ from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( ModelingV2Mode, + NoTarget, resolved_mode, step_contract_enabled, ) @@ -771,16 +772,16 @@ def _validate_runtime(model_config: ModelConfig) -> None: or getattr(sparse, "algorithm", None) != "minimax_m3" or sparse.implementation != "msa" ): - raise ValueError( + raise NoTarget( "this MiniMax-M3 target requires sparse_attention_config.algorithm='minimax_m3' with implementation='msa'" ) if getattr(sparse, "indexer_kv_dtype", None) != "fp8": - raise ValueError( + raise NoTarget( "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") + raise NoTarget("this MiniMax-M3 target requires an FP8 or NVFP4 KV cache") def _sparse_layer_ids(cfg: PretrainedConfig) -> List[int]: @@ -3110,7 +3111,7 @@ def __init__( or getattr(cfg, "bias", False) or getattr(cfg, "tie_word_embeddings", False) ): - raise ValueError( + raise NoTarget( "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" @@ -3156,7 +3157,7 @@ def __init__( sparse_params=None, ) if not self.attn.support_fused_rope() or not self.attn.support_fused_qkv(): - raise ValueError( + raise NoTarget( f"the drafter's attention backend {type(self.attn).__name__} must fuse RoPE and QKV" ) self.attn.update_quant_config(quant_config) @@ -4126,7 +4127,7 @@ def _build_speculation(shell, model_config: ModelConfig): or spec.eagle3_model_arch != "llama3" or getattr(spec, "use_dynamic_tree", False) ): - raise ValueError( + raise NoTarget( "this MiniMax-M3 target speculates with linear one-model Eagle3 (llama3 drafter) only" ) draft_config = ModelConfig.from_pretrained( 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 867001663065..862a48be66f8 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 @@ -16,6 +16,9 @@ The checkpoint declares the vision-language wrapper, and the language model's shape lives in its ``text_config``, so that is what this tree reads. The targets are text-only: they load no vision tower and refuse image input. + +Under ``modeling_v2='default'`` a target that ``route`` returns is used only +where ``enabled_by_default`` allows it: the serving roles it was measured in. """ from __future__ import annotations @@ -25,7 +28,7 @@ from tensorrt_llm.llmapi.llm_args import EagleDecodingConfig from tensorrt_llm.quantization.mode import QuantAlgo -from ..._router_index import NoTarget +from ..._router_index import DeploymentRole, NoTarget # The one GPU architecture these targets are written for: the decode kernels # (the tcgen05 GEMVs, the fused MoE engines, the attention megakernels and the @@ -210,3 +213,23 @@ def route( "tensor-parallel attention and experts, with one-model Eagle3 or no speculative decoding, and an FP8 or " "NVFP4 KV cache" ) + + +# The serving roles each target is enabled in under modeling_v2='default': the deployments it was measured in. +_DEFAULT_ROLES = { + # GB300: the generation side of disaggregated serving. + "ModelingV2MiniMaxM3Nvfp4Sm103Tp4": (DeploymentRole.GENERATION,), + "ModelingV2MiniMaxM3Nvfp4Sm103Tp8": (DeploymentRole.GENERATION,), +} + + +def enabled_by_default(target: type, role: DeploymentRole) -> Optional[str]: + """Why ``target`` stays off under ``modeling_v2='default'`` for an engine in ``role``, or None when it is on. + + ``auto`` and ``require`` still reach a target wherever ``route`` returns it. + """ + roles = _DEFAULT_ROLES.get(target.__name__, ()) + if role in roles: + return None + measured = " or ".join(r.value for r in roles) or "no role" + return f"{role.value} serving; {target.__name__} is enabled by default for {measured} serving" diff --git a/tensorrt_llm/_torch/models/modeling_auto.py b/tensorrt_llm/_torch/models/modeling_auto.py index 0592317d293a..4e8d31648d0f 100644 --- a/tensorrt_llm/_torch/models/modeling_auto.py +++ b/tensorrt_llm/_torch/models/modeling_auto.py @@ -1,7 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from dataclasses import replace from typing import Generic, Optional, Type +from tensorrt_llm.logger import logger + +from .._experimental.modeling_v2._router_index import (ModelingV2Mode, NoTarget, + resolved_mode) from ..model_config import ModelConfig from ..utils import model_extra_attrs from .modeling_utils import (DecoderModelForCausalLM, TConfig, TModel, @@ -74,7 +79,29 @@ def from_config( config.skip_create_weights_in_init = True config._frozen = True extra_attrs = config.extra_attrs + # What the loader put there, before a target that then refuses could + # add to it. + loader_attrs = dict(extra_attrs) with model_extra_attrs(extra_attrs): - model = cls(config) + try: + model = cls(config) + except NoTarget as why: + # A modeling_v2 target refuses, at construction, a runtime that + # routing does not see (set by engine arguments). Under auto + # (and default) that falls back to the built-in implementation, + # as a refusal in routing does; under require it is the error. + target = config.modeling_v2_target + if cls is not target or resolved_mode( + target) is ModelingV2Mode.REQUIRE: + raise + logger.info( + f"modeling_v2: {target.__name__} refused this runtime ({why}); " + "using the built-in implementation") + fallback = replace(config, modeling_v2_target=None) + # replace() starts init=False fields afresh; the built-in model + # reads the loader's settings (GEMM backends, allreduce sizes) + # from extra_attrs. + fallback.extra_attrs.update(loader_attrs) + return AutoModelForCausalLM.from_config(fallback) model.extra_attrs = extra_attrs return model diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index d567596628e1..42f0eb12af48 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -6940,16 +6940,20 @@ def quant_config(self) -> QuantConfig: def quant_config(self, value: QuantConfig): self._quant_config = value - modeling_v2: Literal["off", "auto", "require"] = Field( - default="off", + modeling_v2: Literal["default", "off", "auto", "require"] = Field( + default="default", status="prototype", description= "Whether to build the model from a modeling_v2 target (experimental) " - "instead of the built-in implementation. 'off' never does. 'auto' " + "instead of the built-in implementation. 'default' behaves as 'auto' " + "on the deployments the architecture's routing module enables by " + "default and as 'off' everywhere else. 'off' never does. 'auto' " "uses a target when one claims this deployment and falls back to the " "built-in implementation otherwise. 'require' raises instead of " "falling back.", - json_schema_extra={"type": "Literal['off', 'auto', 'require']"}) + json_schema_extra={ + "type": "Literal['default', 'off', 'auto', 'require']" + }) # TODO: remove backend later backend: Literal["pytorch"] = Field( diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 8b268a9df3f2..fa9fd9d58b9e 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -7099,9 +7099,10 @@ "title": "Model Kwargs" }, "modeling_v2": { - "default": "off", - "description": "Whether to build the model from a modeling_v2 target (experimental) instead of the built-in implementation. 'off' never does. 'auto' uses a target when one claims this deployment and falls back to the built-in implementation otherwise. 'require' raises instead of falling back.", + "default": "default", + "description": "Whether to build the model from a modeling_v2 target (experimental) instead of the built-in implementation. 'default' behaves as 'auto' on the deployments the architecture's routing module enables by default and as 'off' everywhere else. 'off' never does. 'auto' uses a target when one claims this deployment and falls back to the built-in implementation otherwise. 'require' raises instead of falling back.", "enum": [ + "default", "off", "auto", "require" diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index c9a19899dff5..c29fa240e00e 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -2465,9 +2465,10 @@ "title": "Model Kwargs" }, "modeling_v2": { - "default": "off", - "description": "Whether to build the model from a modeling_v2 target (experimental) instead of the built-in implementation. 'off' never does. 'auto' uses a target when one claims this deployment and falls back to the built-in implementation otherwise. 'require' raises instead of falling back.", + "default": "default", + "description": "Whether to build the model from a modeling_v2 target (experimental) instead of the built-in implementation. 'default' behaves as 'auto' on the deployments the architecture's routing module enables by default and as 'off' everywhere else. 'off' never does. 'auto' uses a target when one claims this deployment and falls back to the built-in implementation otherwise. 'require' raises instead of falling back.", "enum": [ + "default", "off", "auto", "require" @@ -8771,9 +8772,10 @@ "title": "Model Kwargs" }, "modeling_v2": { - "default": "off", - "description": "Whether to build the model from a modeling_v2 target (experimental) instead of the built-in implementation. 'off' never does. 'auto' uses a target when one claims this deployment and falls back to the built-in implementation otherwise. 'require' raises instead of falling back.", + "default": "default", + "description": "Whether to build the model from a modeling_v2 target (experimental) instead of the built-in implementation. 'default' behaves as 'auto' on the deployments the architecture's routing module enables by default and as 'off' everywhere else. 'off' never does. 'auto' uses a target when one claims this deployment and falls back to the built-in implementation otherwise. 'require' raises instead of falling back.", "enum": [ + "default", "off", "auto", "require" diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index e58f42598f2d..9e0190a5555a 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -992,6 +992,7 @@ }, { "allowed_values": [ + "default", "off", "auto", "require" 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 d63924daf06b..2369de2065d8 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -32,7 +32,12 @@ from transformers import PretrainedConfig from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + _RESOLVED_MODES, + DISAGG_ROLE_ENV, + DeploymentRole, ModelingV2Mode, + NoTarget, + deployment_role, modeling_v2_resolve, resolved_mode, ) @@ -40,6 +45,7 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_auto import AutoModelForCausalLM from tensorrt_llm.llmapi.llm_args import ( + CacheTransceiverConfig, KvCacheConfig, MTPDecodingConfig, NGramDecodingConfig, @@ -177,7 +183,8 @@ def _qwen4_exp_config(**vision_overrides): @pytest.mark.parametrize("mode", [None, "off"], ids=["default", "off"]) def test_off_resolves_nothing(mode): - """The default must be indistinguishable from modeling_v2 not existing. + """The default must be indistinguishable from modeling_v2 not existing for an architecture whose routing module + enables nothing by default. Leaving the argument alone and an explicit "off" have to behave identically: the common case is that nobody has heard of this package. @@ -451,6 +458,71 @@ def test_an_unknown_mode_is_rejected_when_the_arguments_are_built(): _llm_args("yes") +def test_the_default_mode_is_default(): + assert TorchLlmArgs(model=_DUMMY_MODEL).modeling_v2 == "default" + + +def test_no_target_is_a_value_error(): + """A refusal is the ValueError ``require`` has always raised, so callers that catch ValueError still do.""" + assert issubclass(NoTarget, ValueError) + + +@pytest.mark.parametrize( + ("transceiver", "env", "role"), + [ + (False, None, DeploymentRole.AGGREGATED), + (False, "generation", DeploymentRole.AGGREGATED), + (True, "generation", DeploymentRole.GENERATION), + (True, "gen0", DeploymentRole.GENERATION), + (True, "context", DeploymentRole.CONTEXT), + (True, "ctx1", DeploymentRole.CONTEXT), + (True, None, DeploymentRole.UNKNOWN), + (True, "coordinator", DeploymentRole.UNKNOWN), + ], +) +def test_deployment_role(monkeypatch, transceiver, env, role): + """No cache transceiver is aggregated serving; a disaggregated worker's side is the launcher's role variable.""" + if env is None: + monkeypatch.delenv(DISAGG_ROLE_ENV, raising=False) + else: + monkeypatch.setenv(DISAGG_ROLE_ENV, env) + kwargs = dict(cache_transceiver_config=CacheTransceiverConfig()) if transceiver else {} + assert deployment_role(_llm_args(**kwargs)) is role + + +class _RefusingTarget: + """A target that refuses its runtime at construction.""" + + def __init__(self, config): + raise NoTarget("this runtime is not the one the target is written for") + + +class _BuiltIn: + def __init__(self, config): + self.config = config + + +@pytest.mark.parametrize("mode", ["auto", "require"]) +def test_a_refusal_at_construction_falls_back_unless_required(monkeypatch, mode): + """A target's construction-time refusal builds the built-in implementation under ``auto`` (the mode + ``default`` records) and is the error under ``require``.""" + import tensorrt_llm._torch.models.modeling_auto as modeling_auto + + monkeypatch.setattr(modeling_auto, "get_registered_model_class", lambda arch: _BuiltIn) + monkeypatch.setitem(_RESOLVED_MODES, _RefusingTarget, ModelingV2Mode(mode)) + config = replace(_model_config(_gpt_oss_config()), modeling_v2_target=_RefusingTarget) + config.extra_attrs["nvfp4_gemm_allowed_backends"] = ["cutlass"] + if mode == "require": + with pytest.raises(NoTarget, match="not the one the target is written for"): + AutoModelForCausalLM.from_config(config) + return + model = AutoModelForCausalLM.from_config(config) + assert isinstance(model, _BuiltIn) + assert model.config.modeling_v2_target is None + # The loader's settings reach the built-in model. + assert model.config.extra_attrs["nvfp4_gemm_allowed_backends"] == ["cutlass"] + + # --- speculative decoding is an identity criterion ------------------------------ @@ -781,6 +853,61 @@ def test_minimax_m3_kv_cache_dtype_overrides_the_checkpoint(monkeypatch): assert target.__name__ == "ModelingV2MiniMaxM3Nvfp4Sm103Tp4" +def _m3_default_args(monkeypatch, role, kv_dtype="nvfp4"): + """modeling_v2 left at its default, for an engine serving ``role``.""" + if role == "aggregated": + monkeypatch.delenv(DISAGG_ROLE_ENV, raising=False) + return _llm_args("default", kv_cache_config=dict(dtype=kv_dtype)) + if role is None: + monkeypatch.delenv(DISAGG_ROLE_ENV, raising=False) + else: + monkeypatch.setenv(DISAGG_ROLE_ENV, role) + return _llm_args( + "default", + kv_cache_config=dict(dtype=kv_dtype), + cache_transceiver_config=CacheTransceiverConfig(), + ) + + +@pytest.mark.usefixtures("_minimax_m3_sm") +@pytest.mark.parametrize("tp", [4, 8]) +def test_minimax_m3_default_enables_the_generation_side(monkeypatch, tp): + """With modeling_v2 left alone, the GB300 targets serve the generation side of disaggregated serving, the + deployment they were measured in, and record ``auto`` (so a later refusal falls back).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + config = _m3_resolve_config(_m3_sparse(), None, tp) + target = modeling_v2_resolve(config, _m3_default_args(monkeypatch, "generation")) + assert target.__name__ == f"ModelingV2MiniMaxM3Nvfp4Sm103Tp{tp}" + assert resolved_mode(target) is ModelingV2Mode.AUTO + + +@pytest.mark.usefixtures("_minimax_m3_sm") +@pytest.mark.parametrize( + "role", ["aggregated", "context", None], ids=["aggregated", "context", "unknown"] +) +def test_minimax_m3_default_keeps_other_roles_on_the_builtin_model(monkeypatch, role): + """Aggregated serving, the context side and a disaggregated worker of unknown side build the built-in model + by default; ``auto`` still takes the target there.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + config = _m3_resolve_config(_m3_sparse(), None) + assert modeling_v2_resolve(config, _m3_default_args(monkeypatch, role)) is None + assert modeling_v2_resolve(config, _m3_llm_args("auto", "nvfp4")) is not None + + +@pytest.mark.usefixtures("_minimax_m3_sm") +def test_minimax_m3_default_falls_back_where_route_declines(monkeypatch): + """By default a deployment no target fits (here a BF16 KV cache) builds the built-in model, as under auto.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + config = _m3_resolve_config(_m3_sparse(fuse_qkv_index_projection=False), None) + assert modeling_v2_resolve(config, _m3_default_args(monkeypatch, "generation", "auto")) is None + + # --- qwen4_exp: one-model MTP-Eagle over a linear chain, one rank --------------- _QWEN4_EXP_TARGET = "models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.modeling." diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index d2d1739473ff..97885b124ce8 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -107,8 +107,8 @@ methods: default: pytorch status: deprecated modeling_v2: - annotation: Literal['off', 'auto', 'require'] - default: 'off' + annotation: Literal['default', 'off', 'auto', 'require'] + default: 'default' status: prototype cuda_graph_config: annotation: Union[tensorrt_llm.llmapi.llm_args.DecodeCudaGraphConfig, tensorrt_llm.llmapi.llm_args.EncodeCudaGraphConfig, NoneType] From 6e4caeaba9ef30a4eb9f30656dec9168704e52fb Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 27 Oct 2026 13:08:13 -0700 Subject: [PATCH 13/14] [None][fix] Use TP process groups for MiniMax-M3 workspaces and refresh executor tests Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../ops/m3_mc_allreduce.py | 15 ++-- .../ops/m3_sampled_accept.py | 15 ++-- .../ops/m3_shard_accept.py | 15 ++-- .../ops/m3_mc_allreduce.py | 15 ++-- .../ops/m3_sampled_accept.py | 15 ++-- .../ops/m3_shard_accept.py | 15 ++-- .../distributed/test_mnnvl_workspace_comm.py | 68 +++++++++++++++++++ .../engine/test_decoder_forward_scopes.py | 8 +-- .../engine/test_decoder_runtime_hooks.py | 2 +- .../_torch/executor/test_cuda_graph_runner.py | 1 - 10 files changed, 109 insertions(+), 60 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_mc_allreduce.py index cd1ed18a6cb9..d922b2aadbd2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_mc_allreduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_mc_allreduce.py @@ -89,21 +89,18 @@ class McAllReduceWorkspace: @classmethod def create(cls, mapping, hidden: int, max_tokens: int = 8) -> "McAllReduceWorkspace": from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, _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() + from tensorrt_llm._utils import mpi_disabled + + comm = _get_mnnvl_workspace_comm(mapping) 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() + if not mpi_disabled(): + comm.Free() raise RuntimeError( "McAllReduceWorkspace.create must run outside CUDA-graph capture on every rank" ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_sampled_accept.py index 59c58336c5d7..86b08a0b24fc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_sampled_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_sampled_accept.py @@ -55,20 +55,17 @@ def create(cls, mapping, device) -> "SampledAcceptWorkspace": uc = torch.zeros(words, dtype=torch.int64, device=device) return cls(uc, uc.data_ptr(), 0, 1) from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, _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() + from tensorrt_llm._utils import mpi_disabled + + comm = _get_mnnvl_workspace_comm(mapping) ok = not torch.cuda.is_current_stream_capturing() if not _mnnvl_workspace_all_succeeded(comm, ok): - comm.Free() + if not mpi_disabled(): + comm.Free() raise RuntimeError( "SampledAcceptWorkspace.create must run outside CUDA-graph capture on every rank" ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_shard_accept.py index 2b7d05ad920f..21521591bfa9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_shard_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/m3_shard_accept.py @@ -62,19 +62,16 @@ def create_workspace(mapping, device, slot_rows: int) -> AcceptWorkspace: words = 2 * slot_rows * max(mapping.tp_size, 1) if mapping.tp_size > 1: from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, _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() + from tensorrt_llm._utils import mpi_disabled + + comm = _get_mnnvl_workspace_comm(mapping) if not _mnnvl_workspace_all_succeeded(comm, True): - comm.Free() + if not mpi_disabled(): + 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() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_mc_allreduce.py index 823c78a4ba60..84a9d1e7b0dc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_mc_allreduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_mc_allreduce.py @@ -89,21 +89,18 @@ class McAllReduceWorkspace: @classmethod def create(cls, mapping, hidden: int, max_tokens: int = 8) -> "McAllReduceWorkspace": from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, _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() + from tensorrt_llm._utils import mpi_disabled + + comm = _get_mnnvl_workspace_comm(mapping) 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() + if not mpi_disabled(): + comm.Free() raise RuntimeError( "McAllReduceWorkspace.create must run outside CUDA-graph capture on every rank" ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_sampled_accept.py index 59c58336c5d7..86b08a0b24fc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_sampled_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_sampled_accept.py @@ -55,20 +55,17 @@ def create(cls, mapping, device) -> "SampledAcceptWorkspace": uc = torch.zeros(words, dtype=torch.int64, device=device) return cls(uc, uc.data_ptr(), 0, 1) from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, _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() + from tensorrt_llm._utils import mpi_disabled + + comm = _get_mnnvl_workspace_comm(mapping) ok = not torch.cuda.is_current_stream_capturing() if not _mnnvl_workspace_all_succeeded(comm, ok): - comm.Free() + if not mpi_disabled(): + comm.Free() raise RuntimeError( "SampledAcceptWorkspace.create must run outside CUDA-graph capture on every rank" ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_shard_accept.py index 2b7d05ad920f..21521591bfa9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_shard_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/ops/m3_shard_accept.py @@ -62,19 +62,16 @@ def create_workspace(mapping, device, slot_rows: int) -> AcceptWorkspace: words = 2 * slot_rows * max(mapping.tp_size, 1) if mapping.tp_size > 1: from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, _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() + from tensorrt_llm._utils import mpi_disabled + + comm = _get_mnnvl_workspace_comm(mapping) if not _mnnvl_workspace_all_succeeded(comm, True): - comm.Free() + if not mpi_disabled(): + 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() diff --git a/tests/unittest/_torch/distributed/test_mnnvl_workspace_comm.py b/tests/unittest/_torch/distributed/test_mnnvl_workspace_comm.py index b4314d1657e6..19df230e2390 100644 --- a/tests/unittest/_torch/distributed/test_mnnvl_workspace_comm.py +++ b/tests/unittest/_torch/distributed/test_mnnvl_workspace_comm.py @@ -21,8 +21,13 @@ to make the MNNVL workspace allocation fail and fall back to NCCL. """ +import importlib +from unittest.mock import Mock + import pytest +import torch +import tensorrt_llm._utils as utils from tensorrt_llm._torch.distributed import ops @@ -279,3 +284,66 @@ def test_explicit_request_still_respects_hardware(mnnvl_capable_hardware, monkey mapping = _FakeMapping() mapping.multi_node = True assert not ops.MNNVLAllReduce.is_mnnvl(mapping, torch.int8, explicitly_requested=True) + + +@pytest.mark.parametrize("tp_size", [4, 8]) +@pytest.mark.parametrize("kind", ["mc_allreduce", "shard_accept", "sampled_accept"]) +@pytest.mark.parametrize("setup_succeeds", [True, False]) +def test_minimax_workspace_constructors_use_tp_process_group( + ray_mode, monkeypatch, tp_size: int, kind: str, setup_succeeds: bool +) -> None: + """Run the production constructors without MPI; only GPU allocation is substituted.""" + module = importlib.import_module( + "tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3." + f"minimax_m3_nvfp4__sm_103__tp{tp_size}.ops.m3_{kind}" + ) + pg = _FakeProcessGroup() + pg.GROUP_SIZE = tp_size + pg.backend.group_size = tp_size if setup_succeeds else 0 + mapping = _FakeMapping(pg) + mapping.tp_size = tp_size + no_mpi = Mock(side_effect=AssertionError("MPI must not be used by Ray workers")) + monkeypatch.setattr(ops, "mpi_comm", no_mpi) + monkeypatch.setattr(utils, "mpi_comm", no_mpi) + monkeypatch.setattr(utils, "mpi_disabled", lambda: True) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + synchronize = Mock() + monkeypatch.setattr(torch.cuda, "synchronize", synchronize) + + # CPU buffers exercise tensor initialization and constructor outputs without CUDA. + uc = torch.empty(4096, dtype=torch.int32 if kind == "mc_allreduce" else torch.int64) + mc = torch.empty_like(uc) + handle = Mock(spec_set=["get_uc_buffer", "get_mc_buffer"]) + handle.get_uc_buffer.return_value = uc + handle.get_mc_buffer.return_value = mc + allocate = Mock(return_value=handle) + monkeypatch.setattr(ops, "_make_mnnvl_mcast_buffer", allocate) + if kind == "sampled_accept": + # This host-only C++ op supplies a size; communicator selection stays real. + monkeypatch.setattr(module, "_ws_words_op", lambda world: 4096) + + def create() -> object: + if kind == "mc_allreduce": + return module.McAllReduceWorkspace.create(mapping, hidden=6144) + if kind == "shard_accept": + return module.create_workspace(mapping, torch.device("cpu"), slot_rows=4) + return module.SampledAcceptWorkspace.create(mapping, torch.device("cpu")) + + if not setup_succeeds: + # A borrowed ProcessGroup has no MPI Free(); rejection must preserve the intended error. + with pytest.raises(RuntimeError, match="outside CUDA-graph capture|creation failed"): + create() + allocate.assert_not_called() + synchronize.assert_not_called() + else: + workspace = create() + assert workspace.comm is pg + assert workspace.rank == mapping.tp_rank + assert workspace.world_size == tp_size + assert workspace.handle is handle + allocate.assert_called_once() + assert allocate.call_args.args[0] is pg + assert pg.backend.allreduces == 2 + synchronize.assert_called_once() + no_mpi.assert_not_called() + assert pg.requested_devices == [torch.device("cpu")] * pg.backend.allreduces diff --git a/tests/unittest/_torch/executor/engine/test_decoder_forward_scopes.py b/tests/unittest/_torch/executor/engine/test_decoder_forward_scopes.py index f244495bcf10..d3c72fd0a563 100644 --- a/tests/unittest/_torch/executor/engine/test_decoder_forward_scopes.py +++ b/tests/unittest/_torch/executor/engine/test_decoder_forward_scopes.py @@ -38,7 +38,7 @@ def _runner( ) -> DecoderRunner: # forward reads only the model's extra attributes, the two per-engine objects and the step. runner = object.__new__(DecoderRunner) - runner.model = _ModelWithoutWeights() + runner._model = _ModelWithoutWeights() runner._h2d_cache = skip_cache runner._step_pack = pack runner._forward_decoder = step @@ -175,7 +175,7 @@ class _ScheduledWithContext: def test_greedy_fallback_leaves_the_replayed_graph_outputs_intact(): runner = object.__new__(DecoderRunner) - runner.model = _GreedyFallbackModel() + runner._model = _GreedyFallbackModel() # max_beam_width is read from the runner config; the fallback reads nothing else from it. runner._config = Mock(spec_set=["max_beam_width"], max_beam_width=1) runner.sample_in_graph_callable = None @@ -211,7 +211,7 @@ def test_greedy_only_outputs_follow_the_guided_decoder_switch(keep, monkeypatch) """With an executor-side guided decoder (keep_greedy_only_outputs False) the batch gets logits; otherwise the forward's greedy-only outputs are returned as they are.""" runner = object.__new__(DecoderRunner) - runner.model = _GreedyFallbackModel() + runner._model = _GreedyFallbackModel() runner._config = Mock(spec_set=["max_beam_width"], max_beam_width=1) runner.sample_in_graph_callable = None runner.keep_greedy_only_outputs = keep @@ -231,7 +231,7 @@ def test_greedy_only_outputs_follow_the_guided_decoder_switch(keep, monkeypatch) def test_an_idle_batch_gets_logits_for_its_dummy_row(): """An empty batch (an idle attention-DP rank's dummy row) never keeps greedy-only outputs.""" runner = object.__new__(DecoderRunner) - runner.model = _GreedyFallbackModel() + runner._model = _GreedyFallbackModel() runner._config = Mock(spec_set=["max_beam_width"], max_beam_width=1) runner.sample_in_graph_callable = None hidden = torch.ones(1, 4, device="cuda") diff --git a/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py b/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py index 9fda60feb590..3bfbaee8c913 100644 --- a/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py +++ b/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py @@ -34,7 +34,7 @@ def _runner( ) -> DecoderRunner: # Only the two attributes the threshold reads. runner = object.__new__(DecoderRunner) - runner.sparse_attention_config = sparse_attention_config + runner._sparse_attention_config = sparse_attention_config runner.runtime_hooks = hooks return runner diff --git a/tests/unittest/_torch/executor/test_cuda_graph_runner.py b/tests/unittest/_torch/executor/test_cuda_graph_runner.py index 39b07ef1e76c..57d6cc241592 100644 --- a/tests/unittest/_torch/executor/test_cuda_graph_runner.py +++ b/tests/unittest/_torch/executor/test_cuda_graph_runner.py @@ -225,7 +225,6 @@ def _decoder_runner( spec_config=None, cuda_graph_mem_pool=None, enable_attention_dp=False, - original_max_draft_len=0, original_max_total_draft_tokens=0, is_encoder_decoder=False, mapping=Mapping(), From d1a23f1c10d5d2cb681d4b251ae8a74fb786ac06 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 27 Oct 2026 13:11:12 -0700 Subject: [PATCH 14/14] [None][fix] Keep MiniMax-M3 modeling_v2 opt-in like other models Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../_experimental/modeling_v2/README.md | 3 +- .../modeling_v2/_router_index.py | 58 ------------ .../modeling_v2/models/minimax_m3/routing.py | 24 +---- tensorrt_llm/_torch/models/modeling_auto.py | 2 +- tensorrt_llm/llmapi/llm_args.py | 12 +-- .../schemas/trtllm-serve-config.schema.json | 5 +- .../trtllm-serve-disagg-config.schema.json | 10 +- .../usage/llm_args_golden_manifest.json | 1 - .../distributed/test_mnnvl_workspace_comm.py | 1 + .../modeling_v2/test_modeling_v2_routing.py | 92 +++---------------- .../api_stability/references/llm.yaml | 4 +- 11 files changed, 31 insertions(+), 181 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index d2b5413087f5..f9c8568348f3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md @@ -24,8 +24,7 @@ modeling_v2: require | `modeling_v2` | Behaviour | |---|---| -| `default` (default) | `auto` where the family's routing module enables the target by default (`enabled_by_default(target, role)`, the serving roles it was measured in), `off` everywhere else. A family without `enabled_by_default` is not routed at all. | -| `off` | The resolver returns immediately. Nothing beyond it is imported and behaviour is byte-for-byte what it is today. | +| `off` (default) | The resolver returns immediately. Nothing beyond it is imported and behaviour is byte-for-byte what it is today. | | `auto` | Uses a target when one claims this exact configuration; falls back to the built-in implementation otherwise, logging why. | | `require` | Raises instead of falling back, saying which criterion did not hold. | diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index 99f886fbbbf4..c3fa12b34823 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -42,7 +42,6 @@ import enum import os -import sys from typing import TYPE_CHECKING, Callable, Dict, Optional, Tuple from tensorrt_llm.logger import logger @@ -105,8 +104,6 @@ def step_contract_enabled() -> bool: class ModelingV2Mode(str, enum.Enum): """What to do when a deployment reaches the modeling_v2 resolver.""" - #: ``auto`` on the deployments the routing module enables by default, ``off`` everywhere else. - DEFAULT = "default" OFF = "off" AUTO = "auto" REQUIRE = "require" @@ -127,32 +124,6 @@ class NoTarget(ValueError): """ -class DeploymentRole(str, enum.Enum): - """What this engine serves: both phases, or one side of disaggregated serving.""" - - AGGREGATED = "aggregated" - CONTEXT = "context" - GENERATION = "generation" - #: Disaggregated, but the launcher did not say which side. - UNKNOWN = "unknown" - - -#: Set by ``trtllm-serve`` on disaggregated workers (``context`` or ``generation``; legacy ``ctx*`` / ``gen*``). -DISAGG_ROLE_ENV = "TRTLLM_DISAGG_ROLE" - - -def deployment_role(llm_args: "TorchLlmArgs") -> DeploymentRole: - """The role of this engine, known before the model is built: no cache transceiver means aggregated serving.""" - if llm_args.cache_transceiver_config is None: - return DeploymentRole.AGGREGATED - role = os.environ.get(DISAGG_ROLE_ENV, "").strip().lower() - if role == "context" or role.startswith("ctx"): - return DeploymentRole.CONTEXT - if role == "generation" or role.startswith("gen"): - return DeploymentRole.GENERATION - return DeploymentRole.UNKNOWN - - # The mode each target class was last resolved under (the loader builds the class right after resolving it). _RESOLVED_MODES: Dict[type, ModelingV2Mode] = {} @@ -228,10 +199,6 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti for the architecture, or when that ``route`` raises ``NoTarget``. In ``require`` every one of those but ``off`` raises ``ValueError`` instead. - ``default`` resolves as ``auto`` and then keeps the target only where the - routing module's ``enabled_by_default(target, role)`` returns None (a reason - string declines it). An architecture whose routing module has no - ``enabled_by_default`` is not routed at all under ``default``. """ mode = ModelingV2Mode.of(llm_args) if mode is ModelingV2Mode.OFF: @@ -242,17 +209,6 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti return None arch = pretrained_config.architectures[0] - if mode is ModelingV2Mode.DEFAULT: - route = registered_routes().get(arch) - # The family's ``routing.py`` declares ``enabled_by_default`` next to its ``route``. - enabled_by_default = ( - None - if route is None - else getattr(sys.modules[route.__module__], "enabled_by_default", None) - ) - if enabled_by_default is None: - return None - try: route = registered_routes().get(arch) if route is None: @@ -266,20 +222,6 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti f"modeling_v2: no target fits {arch!r} ({why}); using the built-in implementation" ) return None - if target is not None and mode is ModelingV2Mode.DEFAULT: - role = deployment_role(llm_args) - why = enabled_by_default(target, role) - if why is not None: - logger.info( - f"modeling_v2: {target.__name__} fits {arch!r} but is not enabled by default for this deployment " - f"({why}); using the built-in implementation. Set modeling_v2='auto' to use the target." - ) - return None - logger.info( - f"modeling_v2: {target.__name__} is enabled by default for {arch!r} ({role.value}); " - "set modeling_v2='off' to use the built-in implementation" - ) - mode = ModelingV2Mode.AUTO if target is not None: _RESOLVED_MODES[target] = mode return target 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 54d4575e587c..9b60e91f3165 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 @@ -17,8 +17,6 @@ shape lives in its ``text_config``, so that is what this tree reads. The targets are text-only: they load no vision tower and refuse image input. -Under ``modeling_v2='default'`` a target that ``route`` returns is used only -where ``enabled_by_default`` allows it: the serving roles it was measured in. """ from __future__ import annotations @@ -28,7 +26,7 @@ from tensorrt_llm.llmapi.llm_args import EagleDecodingConfig from tensorrt_llm.quantization.mode import QuantAlgo -from ..._router_index import DeploymentRole, NoTarget, engine_kv_cache_quant_algo, routes +from ..._router_index import NoTarget, engine_kv_cache_quant_algo, routes if TYPE_CHECKING: from tensorrt_llm._torch.model_config import ModelConfig @@ -215,23 +213,3 @@ def route( "tensor-parallel attention and experts, with one-model Eagle3 or no speculative decoding, and an FP8 or " "NVFP4 KV cache" ) - - -# The serving roles each target is enabled in under modeling_v2='default': the deployments it was measured in. -_DEFAULT_ROLES = { - # GB300: the generation side of disaggregated serving. - "ModelingV2MiniMaxM3Nvfp4Sm103Tp4": (DeploymentRole.GENERATION,), - "ModelingV2MiniMaxM3Nvfp4Sm103Tp8": (DeploymentRole.GENERATION,), -} - - -def enabled_by_default(target: type, role: DeploymentRole) -> Optional[str]: - """Why ``target`` stays off under ``modeling_v2='default'`` for an engine in ``role``, or None when it is on. - - ``auto`` and ``require`` still reach a target wherever ``route`` returns it. - """ - roles = _DEFAULT_ROLES.get(target.__name__, ()) - if role in roles: - return None - measured = " or ".join(r.value for r in roles) or "no role" - return f"{role.value} serving; {target.__name__} is enabled by default for {measured} serving" diff --git a/tensorrt_llm/_torch/models/modeling_auto.py b/tensorrt_llm/_torch/models/modeling_auto.py index 4e8d31648d0f..370fa5626590 100644 --- a/tensorrt_llm/_torch/models/modeling_auto.py +++ b/tensorrt_llm/_torch/models/modeling_auto.py @@ -88,7 +88,7 @@ def from_config( except NoTarget as why: # A modeling_v2 target refuses, at construction, a runtime that # routing does not see (set by engine arguments). Under auto - # (and default) that falls back to the built-in implementation, + # that falls back to the built-in implementation, # as a refusal in routing does; under require it is the error. target = config.modeling_v2_target if cls is not target or resolved_mode( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 624a0434114b..e02c42ae24f5 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -7107,20 +7107,16 @@ def quant_config(self) -> QuantConfig: def quant_config(self, value: QuantConfig): self._quant_config = value - modeling_v2: Literal["default", "off", "auto", "require"] = Field( - default="default", + modeling_v2: Literal["off", "auto", "require"] = Field( + default="off", status="prototype", description= "Whether to build the model from a modeling_v2 target (experimental) " - "instead of the built-in implementation. 'default' behaves as 'auto' " - "on the deployments the architecture's routing module enables by " - "default and as 'off' everywhere else. 'off' never does. 'auto' " + "instead of the built-in implementation. 'off' never does. 'auto' " "uses a target when one claims this deployment and falls back to the " "built-in implementation otherwise. 'require' raises instead of " "falling back.", - json_schema_extra={ - "type": "Literal['default', 'off', 'auto', 'require']" - }) + json_schema_extra={"type": "Literal['off', 'auto', 'require']"}) # TODO: remove backend later backend: Literal["pytorch"] = Field( diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index b1de080f73ed..de355f70ce33 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -7187,10 +7187,9 @@ "title": "Model Kwargs" }, "modeling_v2": { - "default": "default", - "description": "Whether to build the model from a modeling_v2 target (experimental) instead of the built-in implementation. 'default' behaves as 'auto' on the deployments the architecture's routing module enables by default and as 'off' everywhere else. 'off' never does. 'auto' uses a target when one claims this deployment and falls back to the built-in implementation otherwise. 'require' raises instead of falling back.", + "default": "off", + "description": "Whether to build the model from a modeling_v2 target (experimental) instead of the built-in implementation. 'off' never does. 'auto' uses a target when one claims this deployment and falls back to the built-in implementation otherwise. 'require' raises instead of falling back.", "enum": [ - "default", "off", "auto", "require" diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 0c968cf32aa8..2bb5e4c910b0 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -2486,10 +2486,9 @@ "title": "Model Kwargs" }, "modeling_v2": { - "default": "default", - "description": "Whether to build the model from a modeling_v2 target (experimental) instead of the built-in implementation. 'default' behaves as 'auto' on the deployments the architecture's routing module enables by default and as 'off' everywhere else. 'off' never does. 'auto' uses a target when one claims this deployment and falls back to the built-in implementation otherwise. 'require' raises instead of falling back.", + "default": "off", + "description": "Whether to build the model from a modeling_v2 target (experimental) instead of the built-in implementation. 'off' never does. 'auto' uses a target when one claims this deployment and falls back to the built-in implementation otherwise. 'require' raises instead of falling back.", "enum": [ - "default", "off", "auto", "require" @@ -8877,10 +8876,9 @@ "title": "Model Kwargs" }, "modeling_v2": { - "default": "default", - "description": "Whether to build the model from a modeling_v2 target (experimental) instead of the built-in implementation. 'default' behaves as 'auto' on the deployments the architecture's routing module enables by default and as 'off' everywhere else. 'off' never does. 'auto' uses a target when one claims this deployment and falls back to the built-in implementation otherwise. 'require' raises instead of falling back.", + "default": "off", + "description": "Whether to build the model from a modeling_v2 target (experimental) instead of the built-in implementation. 'off' never does. 'auto' uses a target when one claims this deployment and falls back to the built-in implementation otherwise. 'require' raises instead of falling back.", "enum": [ - "default", "off", "auto", "require" diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 694c34a4f296..f940f2a80937 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1024,7 +1024,6 @@ }, { "allowed_values": [ - "default", "off", "auto", "require" diff --git a/tests/unittest/_torch/distributed/test_mnnvl_workspace_comm.py b/tests/unittest/_torch/distributed/test_mnnvl_workspace_comm.py index 19df230e2390..c095c20f4064 100644 --- a/tests/unittest/_torch/distributed/test_mnnvl_workspace_comm.py +++ b/tests/unittest/_torch/distributed/test_mnnvl_workspace_comm.py @@ -286,6 +286,7 @@ def test_explicit_request_still_respects_hardware(mnnvl_capable_hardware, monkey assert not ops.MNNVLAllReduce.is_mnnvl(mapping, torch.int8, explicitly_requested=True) +@pytest.mark.cpu_only @pytest.mark.parametrize("tp_size", [4, 8]) @pytest.mark.parametrize("kind", ["mc_allreduce", "shard_accept", "sampled_accept"]) @pytest.mark.parametrize("setup_succeeds", [True, False]) 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 9c02df760b50..cac8aeaedf88 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -33,11 +33,8 @@ from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( _RESOLVED_MODES, - DISAGG_ROLE_ENV, - DeploymentRole, ModelingV2Mode, NoTarget, - deployment_role, modeling_v2_resolve, resolved_mode, ) @@ -186,8 +183,7 @@ def _qwen4_exp_config(**vision_overrides): @pytest.mark.parametrize("mode", [None, "off"], ids=["default", "off"]) def test_off_resolves_nothing(mode): - """The default must be indistinguishable from modeling_v2 not existing for an architecture whose routing module - enables nothing by default. + """The default must be indistinguishable from modeling_v2 not existing. Leaving the argument alone and an explicit "off" have to behave identically: the common case is that nobody has heard of this package. @@ -461,8 +457,8 @@ def test_an_unknown_mode_is_rejected_when_the_arguments_are_built(): _llm_args("yes") -def test_the_default_mode_is_default(): - assert TorchLlmArgs(model=_DUMMY_MODEL).modeling_v2 == "default" +def test_the_default_mode_is_off(): + assert TorchLlmArgs(model=_DUMMY_MODEL).modeling_v2 == "off" def test_no_target_is_a_value_error(): @@ -470,29 +466,6 @@ def test_no_target_is_a_value_error(): assert issubclass(NoTarget, ValueError) -@pytest.mark.parametrize( - ("transceiver", "env", "role"), - [ - (False, None, DeploymentRole.AGGREGATED), - (False, "generation", DeploymentRole.AGGREGATED), - (True, "generation", DeploymentRole.GENERATION), - (True, "gen0", DeploymentRole.GENERATION), - (True, "context", DeploymentRole.CONTEXT), - (True, "ctx1", DeploymentRole.CONTEXT), - (True, None, DeploymentRole.UNKNOWN), - (True, "coordinator", DeploymentRole.UNKNOWN), - ], -) -def test_deployment_role(monkeypatch, transceiver, env, role): - """No cache transceiver is aggregated serving; a disaggregated worker's side is the launcher's role variable.""" - if env is None: - monkeypatch.delenv(DISAGG_ROLE_ENV, raising=False) - else: - monkeypatch.setenv(DISAGG_ROLE_ENV, env) - kwargs = dict(cache_transceiver_config=CacheTransceiverConfig()) if transceiver else {} - assert deployment_role(_llm_args(**kwargs)) is role - - class _RefusingTarget: """A target that refuses its runtime at construction.""" @@ -860,59 +833,24 @@ def test_minimax_m3_kv_cache_dtype_overrides_the_checkpoint(monkeypatch): assert target.__name__ == "ModelingV2MiniMaxM3Nvfp4Sm103Tp4" -def _m3_default_args(monkeypatch, role, kv_dtype="nvfp4"): - """modeling_v2 left at its default, for an engine serving ``role``.""" - if role == "aggregated": - monkeypatch.delenv(DISAGG_ROLE_ENV, raising=False) - return _llm_args("default", kv_cache_config=dict(dtype=kv_dtype)) - if role is None: - monkeypatch.delenv(DISAGG_ROLE_ENV, raising=False) - else: - monkeypatch.setenv(DISAGG_ROLE_ENV, role) - return _llm_args( - "default", - kv_cache_config=dict(dtype=kv_dtype), - cache_transceiver_config=CacheTransceiverConfig(), - ) - - @pytest.mark.usefixtures("_minimax_m3_sm") @pytest.mark.parametrize("tp", [4, 8]) -def test_minimax_m3_default_enables_the_generation_side(monkeypatch, tp): - """With modeling_v2 left alone, the GB300 targets serve the generation side of disaggregated serving, the - deployment they were measured in, and record ``auto`` (so a later refusal falls back).""" +@pytest.mark.parametrize("role", ["aggregated", "context", "generation"]) +def test_minimax_m3_requires_explicit_opt_in(monkeypatch, tp: int, role: str) -> None: + """Every serving role stays on the built-in model unless modeling_v2 is requested.""" from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + monkeypatch.setenv("TRTLLM_DISAGG_ROLE", role) config = _m3_resolve_config(_m3_sparse(), None, tp) - target = modeling_v2_resolve(config, _m3_default_args(monkeypatch, "generation")) - assert target.__name__ == f"ModelingV2MiniMaxM3Nvfp4Sm103Tp{tp}" - assert resolved_mode(target) is ModelingV2Mode.AUTO - - -@pytest.mark.usefixtures("_minimax_m3_sm") -@pytest.mark.parametrize( - "role", ["aggregated", "context", None], ids=["aggregated", "context", "unknown"] -) -def test_minimax_m3_default_keeps_other_roles_on_the_builtin_model(monkeypatch, role): - """Aggregated serving, the context side and a disaggregated worker of unknown side build the built-in model - by default; ``auto`` still takes the target there.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing - - monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) - config = _m3_resolve_config(_m3_sparse(), None) - assert modeling_v2_resolve(config, _m3_default_args(monkeypatch, role)) is None - assert modeling_v2_resolve(config, _m3_llm_args("auto", "nvfp4")) is not None - - -@pytest.mark.usefixtures("_minimax_m3_sm") -def test_minimax_m3_default_falls_back_where_route_declines(monkeypatch): - """By default a deployment no target fits (here a BF16 KV cache) builds the built-in model, as under auto.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing - - monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) - config = _m3_resolve_config(_m3_sparse(fuse_qkv_index_projection=False), None) - assert modeling_v2_resolve(config, _m3_default_args(monkeypatch, "generation", "auto")) is None + kwargs = dict(kv_cache_config=dict(dtype="nvfp4")) + if role != "aggregated": + kwargs["cache_transceiver_config"] = CacheTransceiverConfig() + assert modeling_v2_resolve(config, TorchLlmArgs(model=_DUMMY_MODEL, **kwargs)) is None + for mode in ("auto", "require"): + target = modeling_v2_resolve(config, _llm_args(mode, **kwargs)) + assert target.__name__ == f"ModelingV2MiniMaxM3Nvfp4Sm103Tp{tp}" + assert resolved_mode(target) is ModelingV2Mode(mode) # --- qwen4_exp: one-model MTP-Eagle over a linear chain, one rank --------------- diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 34dc4006e8ce..e3e0e82ddcaa 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -107,8 +107,8 @@ methods: default: pytorch status: deprecated modeling_v2: - annotation: Literal['default', 'off', 'auto', 'require'] - default: 'default' + annotation: Literal['off', 'auto', 'require'] + default: 'off' status: prototype cuda_graph_config: annotation: Union[tensorrt_llm.llmapi.llm_args.DecodeCudaGraphConfig, tensorrt_llm.llmapi.llm_args.EncodeCudaGraphConfig, NoneType]